Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancy
Authors: Cameron Allen
Aaron T. Kirtland
Ruo Yu Tao
Sam Lobel
Daniel Scott
Nicholas Petrocelli
Omer Gottesman
Ronald Parr
Michael Littman
George Konidaris
Venue: NeurIPS 2024 poster (2024)
Keywords: reinforcement learning
partial observability
value estimation
memory
Primary area: reinforcement_learning
TL;DR: Minimizing a discrepancy between different value estimates is beneficial for learning memory under partial observability.

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.

Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancy
Cameron Allen' UC Berkeley?
Aaron Kirtland Brown University
Ruo Yu Tao" Brown University
Sam Lobel Brown University
Daniel Scott Georgia Tech
Nicholas Petrocelli Brown University
Omer Gottesman
Amazon
Ronald Parr Duke University
Michael L. Littman Brown University
George Konidaris Brown University
Abstract
Reinforcement learning algorithms typically rely on the assumption that the envi- ronment dynamics and value function can be expressed in terms of a Markovian state representation. However, when state information is only partially observable, how agent learn such 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 1-discrepancy, is the difference between two distinct temporal dif- ference (TD) value estimates, each computed using TD(A) with different value of A. Since TD(A=0) makes an implicit Markov assumption and TD(==1) does not, discrepancy between these estimates is potential indicator of non-Markovian state representation. Indeed, we prove that the X-discrepancy is exactly zero for all Markov decision processes and almost always non-zero for broad class of partially observable environments. We also demonstrate empirically that, once de- tected, minimizing the x-discrepancy can help with learning 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 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.
Introduction
The dominant modeling frameworks for reinforcement learning [Sutton and Barto, 2018] define environments in terms of 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], its partially observable counterpart, a POMDP
'These authors contributed equally and are ordered alphabetically. Please send any correspondence to camallen@berkele edu, aaron kirt sirtlandebrown edu, and ruoyutao@brovm.co edu. Some of this work was completed while at Brown University.
Work completed while outside of Amazon.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

[Kaelbling et al., 1998]. as long as the underlying state representation contains enough information to satisfy the Markov assumption. The POMDP framework 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, 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 observa- tions 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 pos- sible 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 good state representation for the problem it is solving. The conventional deep-learning-ert wisdom is that the best problem representations come from training a system solve task, rather than from human designers. When faced with a potentially partially observable environment, if we provide the agent with large enough recurrent neural network (RNN), it can perhaps discover an internal state representation "end- to-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 general-purpose purpose and scalable. However, these techniques only implicitly address the problem of learning a Markovian state representation; 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 X-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.. x= = 0), which makes an implicit Markov assumption. and Monte Carlo (MC) estimation (A = 1), which does not, and intermediate À values interpolate between these extremes. By comparing value estimates for two distinct values of X, we can check that the agent's 's observations support Markovian value prediction, and augment them with memory if we find they are incomplete.
Our main contributions*
follows:
1. We introduce and formally define the X-discrepancy and prove that it is exactly zero for MDPs and almost always non-zero for broad class of POMDPs that we characterize. This analysis tells us that our metric reliably detects partial observability.
2. We then consider tabular, proof-of-concept experiment that adjusts the parameters of memory function to minimize the X-discrepancy via gradient descent. For this experiment, we compute the x-discrepancy in closed form. Our results demonstrate that minimizing the X-discrepancy is a viable path to miligating partial observability.
3. Finally, we integrate our approach into deep reinforcement learning agent and evaluate on a set of large and challenging POMDP benchmarks. We find that minimizing the X- 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 single value function.
Overall, we find that the X-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 X-discrepancy remains applicable even without the common assumption that the agent knows the full set of possible POMDP states.
'Code: http:///gtthub.cco/oroom/////// discrepancyyccrcyc
Videos: http:///ammmaaddcceppaaee wttprrreppppcccigt..bbb... .ithub. io

---

2 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 Rt S + F. transition function T A - AS mapping to a distribution over states, discount factor E [0, 11, and initial state distribution po E AS. The agent's goal is to find a policy as : S + AA that selects actions to maximize return, It. the discounted sum of future rewards starting from timestep Ei=0?' I'Tt+i, where is the reward at timestep i. We denote the expectation of these returns as value functions v. (s) Ex [gt St = s) and Q (s,a) a) Ens 19: st = 8, a = a).
t:
9, =
The POMDP framework [Kalbling et al., 1998] additionally includes a set of observations Ω and an observation function + S AΩ that describes the probability ((w)s) of seeing observation w in latent state 8. 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 w does not contain enough information to fully resolve the state s. While agents need not state to behave optimally, they must retain at least enough information across timesteps that the optimal policy becomes clear.
n fully resolve the underlying
We are interested in the learning setting, where the agent has knowledge of the underlying state nor even the set of possible states S (let alone the transition, reward, and observation functions). receives observation Wt E ft 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 E Z is Markovian if any timestep t, the representation žn and action ag together are sufficient statistic for predicting the reward Ti and next representation žς-1, instead of requiring the agent's whole history:
Pr(z Fe/ze, ax) = Pr(ze+1, Fi Fr(2t Ot....: 20, an)
The definition can be applied to states (s), observations (w), and even memory-augmenter obser- vations (defined in Sec. 4), by replacing 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: Vns (s) = Σ s( (R(ss, a) + + EvEves T(s' VT (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,
(((+))
Erg rr + vV(( (8++1) 8: =
which converges to the unique fixed point Vns. A second implication is that the transition and reward functions, and consequently the value functions Vrs and Q have fixed-sized inputs and are therefore easy parameterize, learn, and reuse. Finally, it follows from the Markov property that the optimal policy z can be expressed deterministically and does not require memory [Puterman, 1994].
We can unroll the Bellman equation over multiple timesteps to obtain similar estimator that uses n-step returns: V*s(s) (s) = Ens [9t=t s, = s). where Itst+n := T₁ + 7T++1 + 72r + +... yd y" (Se+n), with Itat-n St if the episode terminates before t + has been reached. The same process works for weighted combinations of such returns, including the exponential average:
V (s) = 1 (1 A)
Σ n=1
gt.ttn
St =
(3)
with ViA= Exs I9t 8 s] as a special case. Equation (3) detines the TD(A) value function as an expectation over the so-called X-return [Sutton, 1988]. Low values of X give more weight to shorter returns where the value term Va (St+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(A) value functions share the same fixed point for any X E [0, 1]: however, if the Markov property does not hold, different may have different TD(A) fixed points. In this work, we seek to characterize this phenomenon and leverage it for detecting and mitigating partial observability.
3

---

Memory
PPO
Action
RNN
TD
A
Ohservation
p.00
CAttor
LPPO
Leasie 8 +
8) (Cyro + Lyme)
T-maze decision process. The agent must remember the initial observa- tion to earn the maximum reward (+1).
Figure I:
Figure 2: X-discrepancy augmented network architec- ture and training objectives.
Detecting Partial Observability
Before introducing our partial observability metric, us first consider the T-maze example of Figure under two candidate decision-make makinn 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 the junction that leads to the positive reward.
To model the T-maze as a POMDP, 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 sampled the start 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 Q as the MDP's "states". This requires new transition and reward functions, Ta and Ra, which we define as: Ta(w' w. a) == Es.'es ((w' s')Ts(s' s, s, Pr(s w) and Ra(w,,) a) = E.es Rs(s, a) Pr(s w), where Pr(s w) is dependett and describes how each hidden state s; E S contributes to the overall environment behavior when we see observation w. 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 w). That particular weighting scheme is the one that averages the dinendepn Pr(st we) over all timesteps, weighted by visitation probability under the agent's policy, and discounted by 7.3 We explain how to construct this expression in Appendix C. We call the MDP defined in this way the effective MDP model for 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 (cqually likely) starting observations. Note that the POMDP model does 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".
principle, an agent could measure partial observability by simultaneously modeling its environment as both 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-lenge length history inputs at each time step, and would come with computational and memory costs that grow exponentially with history length. Instcad, we propose a model-free approach, using value functions conditioned only observations, to approximate this comparison in tractable way.
3.1 Value Function Estimation under Partial Observability
The Bellman equation and its sample-bass 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 (A = 0),
One intuitive way think about Pr(s(w) is that reflects the relative frequency trajectories matching both s and s versus those just matching E- E222 [[[=== 12 mpprsm]]]]

---

where we use the same recurrence relation (2) but now our expectation is sampling from the POMDP:
(w) = Σ
Pr(s| w)
Σ
Faa w) (s((s, +
Σ E
((w') ss')ss(s' a2
.ES
aEA
d'ESwEER
a(a w) -
(Ralw, a) +y
ΣIa(w To(w' (a,w
(4)
EEA
w'ER
where we have suppressed Va*s dependence the observation- 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
(A = 1) does not exploit the Bellman equation at all; simply averages over returns: V>²=
E [et 8, s). Translating the Monte Carlo estimator to POMDPs merely requires one additional
expectation to convert from states to observations:
V 1 (w) E. Lgt Wi = w| Err(
Pr(s| w) Er Ige s) Σ Pr(s W)V (s). SES
(5)
#ES
This means MC (A 1) effectively takes the hidden-state value function Vs of the POMDP and projects it into observation space, whereas TD (A = 0) directly computes the value function for the projected model an MDP, by treating observations as states. Interpolating À between 0 and 1 smoothly varies the value estimate's 's reliance on the effective MDP model.
For generality, we derive an expression (see Appendix C) for Qž-values in terms of 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:
2 ( - -TK,) .1 RAA, where K AII + (1 - x)ØW"
(6)
where the tensor product: contracts two indices instead of one,? with tensors defined as follows:
QA (Ω x A) is a matrix of Q-values;
W (I x S) contains the state-blending weights given by Pr(s w) for observation w;
(S x A x S x AA) is identity tensor with Isas'a' 1[s = s/|1(a a');
T (S x A x S) contains the hidden-state transition probabilities T(s' a,, (,
I (S x S x A) contains the policy spread over hidden states (see below);
& (S x Q) is the observation function, containing probabilities ((w |s);
TT (Ω2 x S x A) combines W with II to obtain probabilities Pr(s, a w):
R (S x A) contains the hidden-state rewards R(s, a).
Let us take a moment to parse this equation. First, 11$ and o Wl are mappings from states to state- action pairs. We call the former the MC policy spread, 11$t _1 1[s s') >., & Twwa which maps states to the expected policy under their observation distribution. We call the latter the TD policy spread, (#W")),,s . En Pa.u Ww.'', 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 X. Multiplying a policy spread tensor on the left by 7 produces an (S x Axxxxxx A) transition-polle tensor, describing the probability of each actton transition (s, a) + A Ass', a') under the policy. Intuitively, Equation (6) says that TD(A) computes Q-values for a POMDP as though state-action pairs are evolving according to TK x = X (1 x)T which is a mixture of the policy dynamics under two transition models: the MC transition-policy (7I1)) and the TD transition-policy (T (T'W" The expression (I- >TK3) R computes the state-space Q-values for this hybrid transition model. Finally, these are projected through W to compute the observation-span Q-values. (See Appendix C for more details.)
3.2 X-Discrepancy
We have shown that, under partial observability, Q value functions may differ for different parameters, due to varying reliance the effective MDP model in the TD(A) estimator. We call this difference the X-discrepancy, and we propose to use it as a measure of partial observability.
'This equivalence justifies our choice Pr(s|w) when defining Th and Ra for the effective MDP model. since states appear in precisely this proportion when we generate experiences in the POMDP. For 3- d-imensional tensors A and B, (AB) s En Aiji Batm, and (A:B), E5e Aij.
5

---

Lambda Discrepaccy (A
No Aliasing
Corridor Aliased
Junction Aliasod
Both Alizsed
0.6
3-
3-
0.4-
Alissing Typs
State
0-
Conridor
Junction
Both
12
0.00
12
12
Oberrvation
0.6
Perfect
Miniig Pacameter (p)
Aliased
Figure 3: T-Maze A-discrepancy, mixing between full and partial observability. (Left) MDP observa- tion function Pperfett- (Right) Various POMDP observation functions 4 Aliased that produce aliased observations the corridor states, junction states, or both. State indices correspond to starting states (0, 1). hallway (2-11) 1). junctions (1 (12, 13), and terminal state (14). Brighter squares indicate higher probability. (Center) X-discrepancy has a minimum at zero for full observability and increases with partial observability. We interpolate between perfect observations and aliased ones, where the obser- vation function is 4 = ppp) p) Perrfect + p & Aliased.
Definition
For a given POMDP model and policy x. the X-discrepancy AA is the weighted norm of the difference between the Q-functions estimated by for {21,
A = 11Q - (UI -
(I- )
The choice of norm be arbitrary, as the norm weighting scheme, long 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.
useful property (that we will prove Theorem 2) is that if the POMDP has Markovian observations, the X-discrepancy is exactly zero for all policies. However, for it to be useful measure of partial observability, we must also show that the x-discrepancy reliably non-zero when observations are non-Markovian. For this, we have the following theorem.
Theorem Given POMDP model P and distinct X + X, if there exists policy Ω AA such that A2." # then A + 0 for all policies except most a set of measure zero.
Proof sketch: We formulate the X-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 says that the X-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 thcorem further suggests that even if particular policy has zero d-discrepancy, small perturbations to that policy will almost surely detect partial observability it is present.
We also find that increasing amounts of partial observability lead to larger X-discrepancy. Figure 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 y Rs 1 and use the norm (see Appendix G.2) The results confirm that the X-discrepancy is useful indicator of partial observability, provided that it is non-zero for at least one policy.
maximum
3.3 What conditions cause the X-discrepancy to be zero?
We can characterize the cases in which the X-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 (KA x²x) is zero (which we will show implies Markovian observations), or
Such artifacts concern in normal environments, since the observation function is fixed.

---

lanctions
Parity Check with Random Memory
Normalized Return
100
400
200
0
0.0
10 10-1
10-3
A-discrepancy (A)
Memoor Size (bits)
Figure 4: (Left) The Parity Check environment, a POMDP with zero X-diserepancy for every policy,
(Center) Almost any randomly -iittmi-iiiiilizeed 1-bit memory function reveals a X-discrepancy. (Right) Memory optimization increases normalized return of subsequent policy learning, whereas memoryless policy optimization fails to beat the uniform random baseline.
is non-zero but is projected away by T, RSA, and/or W. We first consider when the policy spread which are the only terms that depend on X are equal, i... K K2²
Theorem For any POMDP P and any A1 + A2, K Kλa 22 if and only if is a block MDP.
Proof sketch: Using Eq. (6), K K2² can be simplified to IT ΦW which is satisfied if and only if &W = T, i.e. each observation is generated by exactly one hidden state. Proof in Appendix
Thus, when the policy spread tensors equal, it means there is no state aliasing and P is block MDP. Now consider the cases where the difference between K is projected away by T. RSA,, and/or W. starting from the policy spread tensors in Definition 1 and working
way outwards.
the policy spread tensors differ, the transition tensor may project these differences away (TK²) TKA2), for example, if the transition probabilities leading 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-policess may be projected away by the state- blending weights W. the reward function RSA, or both. ([[I-TKK))) RSA = W(I ([I-~T7 yT > > - RSA.A In this case, the effective MDP may be lossy, which is a limitation of the X-discrepancy metric. Fortunately, show in the example below.
such environments appear to be rare, and they are easy handle,
Parity Check Environment. The Parity Check POMDP of Figure has four equally likely starting states, each associated with 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- perfectly symmetric with respect rewards and observations. 9 Every observation has zero expected return, and so the X-discrepancy is zero for all policies. Our analysis in Appendix H.5 suggests such edge-cases are exception and not the rule. Even so, such examples are easy to handle: we can simply add small amount of memory. If the environment is truly an MDP, adding memory should have no effect, but if it a POMDP, a randomly initialized memory function (defined in the next section) may be enough break symmetry and a
for the first eight states differ with and without aliasing, but the differences are
a X-discrepancy (as seen in Figure 4, center).
Taken together, Theorems and 2 suggest that the X-discrepancy could a useful measure for detecting partial observability. In the next section, we demonstrate the efficacy of using the X- discrepancy to learn memory functions that mitigate partial observability.
4 Memory Learning with the X-Discrepancy
The X-discrepancy can identify that we need memory, but it cannot yet tell us what to remember. For that, we must replace the POMDP in Definition with memory-augmentee version. In general,
This environment requires projection both W and SAA to zero the differences; neither alone sufficient.
7

---

Performance with Memory Optimization
Normelized Return
1.0
Memory Size (bils)
0.8
0
0.6
2
33
0.4
T-maze
Paint
Network
Tiger
4xt3
Shuttle
Cheese
Figure 5: Memory optimization increases normalized return of subsequent policy gradient learning. Performance is calculated as the expected start-state value, and is normalized between 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.
memory can be any mapping from agent histories (Wo,00, 00, 1t-1,, w:) 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 μ x A M + .M as a mapping from an observation, action, and memory state, (si, (4i,a, m)) to a next memory state m' = suww,, a, m)) Given a POMDP P. memory function jJ induces a memorry-uugmented POMDP pH, with extended states SM = S x M, actions AM A x M, and observations 12M f x M. The augmented transition dynamics TM preserve the original transition dynamics for states while allowing the agent full observability and control over memory states M. Memory functions naturally lead to memory-augmeent
policies TT f2 + AAM M and value functions V= 22 + R and Qr, Ω x AM - TR that reflect Appendix G.1.
the expected return under such policies. For For details,
The x-discrepancy (Definition 1) applies equally well be used to determine whether particular memory function μ induces POMDP PH with Markovian observation space S.M. We can also use as a training objective for learning such a memory function. We conduct proof-of-concept experiment on several classic POMDPs for which we can obtain closed-form gradients of the X-discrepancy with respect to a parametrized memory function. In each domain, randomly generate set of stochastic policies, select the one with maximal X- discrepancy A. and adjust the parameters of memory function μ to minimize A via gradient descent. Figure shows the improvement in policy gradient performance due to the resulting memory function for increasing memory sizes. The details of this experiment provided in Appendix H. We also run the same experiment on the Parity Check example, and provide results in Figure 4 (right).
semorr--augmentn POMDPs, and can thus
We see that the X-discrepancy can help mitigate partial observability, but is somewhat inconveniently defined in terms of the closed-form value function fixed-points. Fortunately, with the appropriate choice of weighted norm, pairs (or, in the case memory, observation-menttn action tuples):
estimate the x-discrepancy purely from sampled observation-action
|| = (B( (Q (w, -2 (w, 1/2
A A2 = 1 - Q²a11
(7)
The norm is taken over the on-policy joint observation-action distribution (r(w) which allows us 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 X-discrepancy detect partial observability in theory and mitigate under certain idealized conditions. Now we demonstrate how to integrate our metric into sample- based deep reinforcement learning to solve problems requiring large, complex memory functions.
5.1 Combining the X-Discrepancy with PPO
To minimize the X-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 our base algorithm due to its strong performance many POMDPs [Ni al., 2022], and since the discrepancy is natural extension of generalized advantage estimation [Schulman et al., 2016), which
8

---

RockSample 11 11
RockSample 15 15
RNN+ LD
Episode Reward
20-
0.2
0.6
1.0
Timestep
10P
Timestep
x10
Battleship 10 x 10
Partially Observable PacMan
Episnde Reward
15-
2000
PINGLD
10-
1000
FRO+RNN
5-
0.0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
L.
Tincstep
x107
1 Timestep
X10
RNN
Figure 6: (Left) The X-discrepancy auxiliary objective (LD) improves performance over recurrent (RNN) and memoryless PPO. Learning curves shown are the mean and 95% confidence interval over 30 runs. (Right) PacMan memory visualization. The agent moves within the maze (middle), and we reconstruct the dot locations from the agent's 's memory. RNNs (bottom) benefit from the X- discrepancy auxiliary loss (LD, top).
is used in PPO. In this algorithm, recurrent neural network [Amari, 1972] (specifically gated recurrent unit, or GRU [Cho et al., 2014J) 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 critic. The critic is usually a value function network which learns a truncated TD(A) value estimate for its advantage estimate.
To estimate the X-discrepancy, we learn two TD(A) value function networks with different X, parame- terized by 0y, and 0y,2 respectively, and minimize their mean squared difference auxiliary loss:
(zx) (0) E. (Vor (x))
oy.a (z+) )
2
8
where 2, = PilaNN 2;-)) 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 X-discrepancy. However, the theoretical results in Section only hold for value function fixed points and not necessarily for estimated values prior to convergence. Inaccurate value functions may lead to inaccurate X- discrepancy. Fortunately, value functions often provide a useful learning signal before convergence, and we observe in the following sections that the X-discrepancy is also useful during learning.
5.2 Large Partially Observable Environments
We evaluate approach 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, /1) 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. See Appendix I.2 for more details.
Wheee domains have extremely large state spaces: Battleship has ~2 states, without accounting for possible ship locations, and PacMan has ~ 5,0 states, even when ignoring food dot configurations.

---

5.3 Experiments
We conduct experiments with regular PPO, recurrent PPO, and our --discrepacc--augneee rent PPO in Figure 6. We also visualize the agent's 's memory for the partially observable PacMan environment by reconstructing the dot locations from the RNN latent state show where the agent "thinks" it has been (see Appendix L.5). We performed a hyperparameter sweep for each method and report learning curves for undiscounted return (see Appendix 1.3 for discounted learning curves and Appendix 1.4 for additional experimental details).
recur-
The X-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 X-discrepancy helps the agent learn the missing features that allow 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 I6).
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 X-discrepancy. Indeed, we find that a large difference between A1 and Xę was beneficial, with optimal As close to either 0 or In addition to this, our hyperparameter sweep also includes a PPO variant with two different TD(A) value functions but without the X-discrepancy auxiliary loss described Section 5.1. These agents were never selected in the sweep; using the loss in Equation 8 seems to only help with performance in tested environments.
6 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 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 rep- resentations [McCallum, 1996] or state controllers [Meuleau et al., 1999]. Modern approaches mostly use RNNs [Amari, 1972] trained via backpropagation through time (BPTT) [Mozer, 1995] to tackle non MarkMvanov decision processes [Schmidhuber, 1990]. Various approaches use recurrent function approximation to learn better state representations. One successful approach is learning a re- current 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 com- pare 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 recurrent dynamics model [Hafner et al., 2020] to facilitate planning alongside reinforcement learning. These approaches learn their represen- tations implicitly to improve prediction error, rather than explicitly to mitigate partial observability.
7 Conclusion
We introduce the X-discrepancy: an observable and differentiable measure of non-Markovianity suitable for mitigating partial observability. The X-discrepancy is the difference between two distinct value functions estimated using TD(A), for two different values. We characterize the X-discrepancy and prove that it reliably distinguishes MDPs from POMDPs. We then use it as memory learning objective and demonstrate that minimizing it in dloeddform helps learn useful memory functions in small-scale POMDPs. Finally, we propose deep reinforcement learning algorithm which leverages the x-discrepancy as an auxiliary loss, and show that significantly improves the performance of baseline recurrent PPO agent set of large and challenging partially observable tasks.
10

---

Author Contributions
CA. AK and RYT led the project. CA, OG, GK, MLL and SL came up with the initial idea. CA, OG, MLL, SL, and DS conducted the first conceptual investigations and proof-of-coneept experiments. AK led the theoretical work, with support from CA, SL, RP, and RYT. RYT the algorithm development, with support from CA, SL, and GK. CA and RYT led the implementation and experiments, with support from SL, NP, and DS. RP discovered the class of parity check examples and showed that they have zero x-discrepancy. CA led the writing, with support from AK and RYT. OG, GK, MLL, SL, and RP advised on the project and provided regular feedback.
Acknowledgments and Disclosure of Funding
Many thanks to Saket Tiwari, Anita de Mello Koch, Sam Musker, Brad Knox, Michael Dennis, Stuart Russell, and our colleagues at Brown University and UC Berkeley for their valuable advice and discussions towards completing this work. Additional thanks to our reviewers for comments on earlier drafts.
This work was generously supported under NSF grant 1955361 and CAREER grant 1844960 to George Konidaris, NSF fellowships Aaron Kirtland and Sam Lobel, ONR grant N00014-22-1 2592, a gift from Open Philanthropy to the Center for Human-Compatibl AI2050 Senior Fellowship for Stuart Russell from the Schmidt Fund for Strategic Innovation.
AI at Berkeley, and an
References
Cameron Allen, Neev Parikh, Omer Gottesman, George Konidaris. Learning Markov state abstractions for deep reinforcement lcarning. Advances Neural Information Processing Systems, volume 34, pages 8229-8241, 2021.
Shun-ichi Amari. Learning patterns and pattern sequences by self-organizing nets of threshold elements. JEEE Transactions on Computers, (-21(11::1197-2200 1972.
Bram Bakker. Reinforcement learning with long short-term memory. In Advances in Neural Information Processing Systems, volume 14, 2001.
Mare G. Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C. Machado, Subhodeep Moitra, Sameera Ponda, and Ziyu Wang. Autonomous Navigation of Stratospheric Balloons using Reinforcement Learning. Nature, 588(7836)) (786))77777711 (1336)777-82277 82, 2020.
Clément Bonnet, Daniel Luo. Donal John Bymee Shikha Surana, Sasha Abramowitz. Paul Duckworth, Vincent Coyette, Laurence Illing Midgley, Elshadai Tegegn, Tristan Kalloniatis, Omayma Mahjoub, Matthew Mac- farlane, Andries Petrus Smit, Nathan Grinsztajn, Raphael Boige. Cemlyn Neil Waters, Mohamed Ali Ali Mimouni, UIrich Armel Mbou Sob, Ruan John de Kock, Siddarth Singh, Daniel parelos-Blanco, Victor Lc, Amu Pretorius, and Alexandre Laterre. Jumanji: diverse suite of scalable reinforcement lcarning environ- ments JAX. International Conference on Learning Representations, 2024.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary. Dougal Maclaurin. George Necula. Adam Paszke, Jake VanderPlas, Skye Mondderman-Milnee and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
Anthony R. Cassandra. The POMDP page. https://www pomdp. ceomdp.ogg/........ 2003.
Anthony R. Cassandra, Leslic P. Kaelbling, and Michael L. Littman. Acting optimally in partially obscrvable stochastic domains. In Proceedings of the Twelfth AAAI Conference on Artificial Intelligence, pages 1023- 1028, 1994.
Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In Alessandro Moschitti, Bo Pang. and Walter Daelemans, editors, Proceedings of the 2014 Conference on Empirical Methods Natural Language Processing (EMNLP), pages 724-1733, 2014.
Lonnie Chrisman. Reinforcement learning with perceptual aliasing: The perceptual distinctions approach. In Proceedings of the Tenth AAAI Conference Artificial Intelligence, pages 183-188, 1992.
Shi Dong. Benjamin Van Roy, and Zhengyuan Zhou. Simple agent, complex environment: Efficient reinforce- ment learning with agent states. Journal of Machine Learning Research, 23:1-54, 2022.
11

---

Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient with rich observations via latent state decoding. In International Conference on Machine Learning. pages 1663-1674. Proceedings Machine Learning Research, 2019.
Antoine Galataud. HVAC processes control: Can Al help?. 2021.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference Learning Representations, 2020.
Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable MDPs. In Proceedings of the 2015 American Association for Artificial Intelligence, 2015.
Nicolas Heess, Jonathan J. Hunt, Timothy P. Lillicrap, and David Silver. Memory-based control with recurrent neural networks. Deep Reinforcement Learning Workshop at NeurlPS, 2015.
Herbert Jaeger. Tutorial on training recurrent neural networks, covering bppt, rtll, ekf and the echo state network approach. Porsunsrschu Informattionseechmik,, 2002., 5, 2002.
Leslie P. Kaelbling, Michael I.. Littman, and Anthony R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial Intelligence, 101(1-2):99-334 1998.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Steven G. Krantz and Harold R. Parks. primer of real functionn. Springer Science & Business Media, 2002.
Nicholas Kushmerick, Steve Hanks, and Daniel S. Weld. An algorithm for probabilistic planning. Artificial Intelligence, (0(1-)):399-286 1995.
Lihong Li, Thomas J. Walsh, and Michael L. Littman. Towards unified theory of state abstraction for MDPs. In Proceedings of the Ninth International Symposium on Artificial Intelligence and Mathematics, 2006.
Long-Ji Ji Lin and Tom M. Mitchell. Reinforcement learning with hidden states. From animals animats, 2: 271-280, 1993.
Michael L. Lituman. An optimization- basse categorization of reinforcement learning environments. In From Animals to Animats Proceedings of the Second International Conference on Simulation of Adaptive Behavior. The MIT Press, 1993. ISBN 9780262287159, 1599
Qinghua Liu, Alan Chung. Csaba Szepesvari, Chi Jin. When is partially observable reinforvement learning not scary? In Proceedings of Thirty Fifth Conference on Learning Theary, volume 178 of Proceedings of Machine Learning Research, pages 3175-5220, 2022.
Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems.
2022.
Andrew Kachites McCallum. Reinforcement learning with selective perception and hidden state. University of
Rochester, 1996.
Nicolas Meuleau, Leonid Peshkin, Kee-Eung Kim, and Leslie P. Kaelbling. Learning finite-state controllers for partially observable environments. In Proceedings of the 15th Conference on Uncertainty in Artificial Intelligence, page 427-436. 1999.
Boris Samuilovich Mityagin. The zero set of real analytic function. Matematicheskie Zametki, 107(3))473- 475. 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu. Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518:529 113,, 2015.
Volodymyr Maih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings of The 33rd International Conference on Machine Learning volume 48 of Proceedings of Machine Learning Research, pages 1928-192 1937, 2016.
Michael Mozer. A focused backpropagation algorithm for temporal pattern recognition. Complex Systems. 3. 1995.
12

---

Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th Interational Conference Interaational Conference Machine Learning, page 807-814, 2010.
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinev. Recurrent model-free RI. can be strong haseline for many POMDPs. In Proceedings of the 39th Interaational Conference Machine Learning, volume 162, pages 16691-16723, 2022.
Christos H. Papadimitriou and John N. Tsitsiklis. The complexity of Markov decision processes. Mathematics of Operations Research, 12(3):41 450, 1987.
Ronald Parr and Stuart Russell. Approximating optimal policies for partially observable stochastic domains. In Pruceedings of the Fourteenth International Joint Conference on Artificial Intelligence Volume 2, page 1098-1094, 1995.
Razvan Pascanu, Tomas Mikolov, Yoshua Bengio. On the difficulty of training recurrent neural networks. Proceedings of the 30th International Conference on Machine Learning, volume 28, pages 1310-1318, 2013.
Martin Puterman. Markov Decision Processes. Wiley, 1994.
Jürgen Schmidhuher. Reinforcement learning in Markovian and non-Markovian environments. In Advances Neural Information Processing Systems, volume 1990.
John Schulman, Philipp Moritz, Sergey Levine, Michael L. Jordan, and Picter Abbeel. High-dimensional continu- ous control using generalized advantage estimation. In International Conference on Learning Representations, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, 2017.
David Silver and Joel Veness. Monte-Carlo planning in large POMDPs. In Advances Neural Information Processing Systems, 2010.
Trey Smith and Reid Simmons. Heuristic search value iteration for POMDPs. Proceedings of the 20th Conference on Uncertainty in Artificial Intelligence, page 520-527, 2004.
Richard Sutton. Learning predict by the methods of temporal differences. Machine learning, 3(1):9444, 1988.
Richard Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. Bradford Book, 2018.
Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforeement learning with function approximation. In Advances in Neural Information Processing Systems, volume 12, pages 1057-1063, 1999.
Ruo Tao, Adam White, and Marlos C. Machado, Agent-state construction with auxiliary inputs. Transactions on Machine Learning Research, 2023.
Daan Wierstra, Alexander Foerster, Jan Peters, and Jürgen Schmidhuber. Solving deep memory POMDPs with recurrent policy gradients. In Artificial Neural Networks ICANN 2007, pages 697-706, 2007.
Zongzhang Zhang, Michael L.. Littman, and Xiaoping Chen. Covering number as a complexity measure for POMDP planning and learning. In Proceedings of the Twenty-Sixth AAAI Conference on Artificcal Intelligence, 2012.
13

---

A Limitations
The main theoretical limitations of our work come from the class of POMDPs that the X-discrepancy is relevant to. The X-discrepancy is capable of detecting non-Markovian observations in a "large" class of POMDPs, which have described using an analytic almost-all characterization, However, there do exist POMDPs in which the fact that observations are Markorann cannot be detected by X-discrepancy, and we provide one such example Section 3.2. The examples such POMDPs we have found seem to exist on "knife's s edge", meaning that slightly perturbing the parameters of the POMDP results in a non-zero X-discrepancy. Furthermore, incorporating memory into value function seems to reliably produce x-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 A-discrepancy. The auxiliary loss in Section only an approximation of the X- discrepancy. It minimizes the difference in value function estimates opposed to the difference in fixed points that we consider in Section The loss function is therefore departure from our theoretical analysis, though the empirical results suggest that the value functions need not have fully converged for the discrepancy to useful. Future work ought to investigate why this works so well in practice.
B Broader Impact
As work studying the foundations of partially observable reinforcement learning, this paper does not have any partieular 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.
14

---

C TD(A) Fixed Point
Here we derive the fixed point of the TD(A) action-value update rule POMDP, following the Markov version by Sutton [1988]. First, define the expected return given initial observation wi and initial action ag as
En (G' Wo. ad)
Pr(solwo)
Pr(ro so: ro
Pr(so wo)
Pr(si150, ao) Σ Σ
Pr(wi[si)Pr(a (w)))//))) Σ Pr(mils,,air Pr(wi |s Pr)( Pr(sssss)))
Pr(solwb wo)
Σ
Pr(si 80: ao) >
Pr(wp) $2 Pr(ag]w2)
Pr(r2|s2, 02)r2
+
We can define the n-step bootstrapped update rule from this given a value matrix Qx by replacing part of the term with coefficient 7" with Qx value, e.g. for the 2 case, get
Q (wo, ao) = Σ Pr(sol wo
Pr(rolso, ro
Pr(sol40) Σ
Pr(si]s0: ao) Σ Σ Pr(w((s))PP((((uu))
Σ Pr(rils1, a1)r1
Pr(selui) Σ Pr(silsus ao) E. Σ, a1
Pr(w118111 (((pall))p((auuu))
8p
Pr(62/01, a1 E w2 E, Pr(uz] Pr((2)2 (4) u2).
Translating these expressions into matrix notation, we have
E, (G" ao) =
Σ 5o
W Wwp.eo Rep
W Wap.ss Σg
T Ts0.a0.51 Σ wi
Σ $ Psy,uy Tay Rey,0s
81
+
W Wa
Σ, 81
T Top. Σ wl
Σ
Tw ssy
Tsy.a1.32
80
82
Saz Twy Razaa
W2
02
where the terms W, T, R, and T are in Equation 6, and is the Ω - AA policy written as an Ω x A tensor with entries in [0, 1]. In particular, W... = Pr(s|w), which averages Pr(s:/wn) all timesteps, weighted by visitation probability and discounted by 7. This well-defined stationary quantity, and it can be computed as follows. First solve the system Cr = to find the discounted state occupancy counts
c(s), where = Po is the initial state distribution over and CI accounts for the policy-de state-to-state transition dynamics T defined by T = Eu Ea Pr(w]s) Pr(a|w) Pr(s']s, (ss a). Then Pr(s(w)) o c(8)* (w)(s) so we can just multiply these terms together and renormalize. For the 2-step bootstrapped update rule above, we have:
Q6 (1+1) (wo. ao)
Σ
W wo Reg.oo
W 5o E, T :p.
Σ
Σ
T s1 Twr.as aa Ray.
W Wwn.sp
Σ,
T Ton
T Ponc Twr.a T an .82
80
52
*
T Paziwa
I (10) (W2, a2).
15

---

Let II be an fxxΩxx representation of the Qx policy with Hww' = 1| = wT Likewise, let 11s be the effective policy over latent states, an SxSxAA representation of the matrix 2T,, i.c. I 2. 1/s 1 ( )...) Given matrix of action-values Q the n-step update rule is:
Q = Qr,n (Q = W
where RSA is the matrix of reward values described in Equation 6, and denote double tensor contraction, while all other tensor products contract a single index.
We also have the standard definition of the TD(A) update rule as:
Q f++1) = - A) (1- n=1 X 1..
We concerned with the fixed point of this update rule, which refer to Q
2 = (1- (1-1) x" 2a=
( (
Separating this into a reward part with factor RSA and value part with factor Q we find that the value part is
(1 A)W ( DR (AYTI1S)a=)) "7700
(1- -
and for the reward part, we have the coefficients R.A.A. in the table below for values of n and
k
0
2
n1
2
X
XyTII
3
x
X-TIS
where each term is multiplied by (1 X)W in front. We can then see by summing over rows before columns that the reward part is:
(1 (1- Σ k=0 T= (
W - -
So we rewrite Q as follows:
Q: = (1SSA-2711)) -1 : + + (11-ATTTI:QQ.))
Now let W = 1 which is 12 Ax x S x A, and T IA:A which is Ω S A. Here, @ means the Kronecker product. This essentially repeats the W matrix |A] times to incorporate WG. This is because
actions into the tensor. Note that for any Sx.A tensor G, (WA:G) (W- (WA:G) E.ar W.A wase Gsa', and the only nonzero terms this sum are those such that a'. For these indices, W.A wasa' Wws, SO Es,a A Wasa Goa
Σ, Goa (WG)
Also, let F = (JSA - XyTTIT) -1 which is S Ax S x A. Then we find:
WA: Q² =>==A: ( (1--N7TIT)) -1 (RSA + (1 x)TTTI:
-W-A(F:(RRAA + (1-11111)))) -1)x)TTTT1:)
-WA:P:RSA + WA:F 1)
16

---

At this point, we can subtract the second tern on the right hand side from both sides, factor out Q on the right, and multiply by (IRA (1 - F.Tm on the left of both sides to obtain:
2 (1 - (1 >) -1 WA:F:
WA:(11A - 1)1:-Q11:WA> -1 FFRRA
- :.P:
PR:RAA
(1 (1NN)))-TTT
+ +1- -1- (I-)))
where the last equality follows from expanding the geometric series. Now we use the identity
(A B) 2k=o(A "B)= to find:
2² w (F-1-(1- -1 (1- >>>TTT"">> RSA
W (I ~I (AII + -)))w")) RAS
To recap our previous definitions, W an ΩXxS tensor, I is an SxAxSxx tensor, T is an AxAxS tensor, II is an SxSx.A tensor, T is an SxΩ tensor, W is an fx5xx tensor, and RSA is an SxAA tensor.
Lastly, we briefly note that one get the V values by replacing W in the above equation with Will and using a double contraction:
V = -.: 7T (AII +(1 (1 - RSA,
We can confirm that V2 En Tw,n Qwaas by rewriting the expression on the right as follows:
Σ Tur Qwia
Σ Tusa
Σ
Wa
aa
Ba.a'
e,a!
Σ Tura
WA
W.nasaa Bsaa
a
Σ Tusa Σ
Ww. Br,a
Σ Tu,a Wo.e Bea
VA
where B == -1- SIK +(1 )) -1 RSA.A is an S x A tensor.
17

---

D Proof of Theorem 1 (Almost All)
In this section we prove Theorem that there is either a A-discrepancy for almost all policies for no policies. Fix À and X. Recall that we define the X-discrepancy as follows:
Ax 2 - Q 2.T - - (1-
:RSA
where K = -lk + (1 - xlow" Let Y be the largest open set in the space of stochastic Q2 x A matrices, considered as subset of IR (1), Now consider the X-discrepancy as a function of the policy In other words, we define
A Y - IR π - Q - Q
Let X be open subset of IR" We say that function X + R is real analytic X if for all x E X. f can be written as convergent power series 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 two analytic functions is analytic where the denominator nonzero,
3. a real analytic function on domain X is either identically zero or only zero on set of measure zero [Mityagin, 2020].
We will also use the fact that for an invertible matrix A, each entry Anj is analytic in the entries of A where the entries of A yield a nonzero determinant. We can prove this by first writing A 1 det(A) aAj(A) (et( "A) where adj A is the adjugate of A and cof ([A) is the cofactor matrix of A. Each entry of the cofactor matrix cofactor that 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
Next, we will show that is an analytic function. Note that the variable terms in the equation are W, W.l, S T. I1S and 4. Of these, SASA, T. and & are constant with respect [ so each entry of $1$ is analytic on Y in the entries of т. Likewise, P.j Eks Duo Tiaj is analytic Y. Therefore, the state-occupancy counts (T > 'po where Po contains the initial state probabilities, are the composition of analytic functions and thus analytic on Y. W. Σ. F,,
11 En 1[i
Po + TPT Po + (PT +...
is analytic on Y for the same reason, and the
denominator of Wws. E. Co' is nonzero for observations able to be observed with nonzero probability. Finally, A is then composition of analytic functions on Y and thus analytic itself.
To handle the norm weighting, we note that w, is analytic in π as wr (1, (L, (alw)), and the dot product of ww with A is also analytic. Now, we use the fact mentioned above that the zero set of nontrivial analytic function is of measure zero. Therefore, the zero set of A W, either zero for all policies or zero only on set of measure zero. To finish, we note that because norms are positive definite, A if and only w = 0, so this result extends the normed X-discrepancy as well.
18

---

E Proof of Theorem 2 (Block MDP)
In this section, we prove Theorem 2 concerning when the system is a block MDP. Recall that in Equation we define K² XII + (1 >) Suppose K K We can rewrite this as (A X')1I= - x> (X- X'( - 0. This implies that either X = X or II. = DW"
Recall the definition of I and Will in Section C. Expanding the equation ITo - using these definitions, we have that
II. 3, 1|s==|1pp =') -[s-= s') Σ Pr(w|s)Pr((ak))
($W") 1,,0 De.. w.s'
w. Σ Pr(ulss)Pr(al))PPP(((())
Then by setting the rightmost expression in each equation equal and simplifying with the indicator function, have that for all i,
E Prollss)P((aaa))PPrr f (E,Pr((ss))P(aapp
i= i i j.
We will first consider the case where + j. We have that for all i + j and all k. _.. Pr(w]s;; Pr(ax]() (r(aa)l))P(ssllaa)) = 0. Because each term in the sum is nonnegative, this is equiva- lent to the statement that for all i + j. all k. and all w, Pr(w[s; (r(Wss:)P((as) Pr(as Pr(s)() = 0. Now note that for all observations w. there exists some k' such that Pr(ax (w) > 0. Therefore, we have that for all * j and all w, there exists k such that Pr(ag/ (w) > 0. This implies that for
/ j and O. This means that if state 8; produces an observation w, then w cannot be produced by any other reachable state 8, * 8,, where two states are said to be reachable if there exists sequence of actions sampled from the policy that enable the agent to reach state s; from 8j with nonzero probability. In other words, each observation uniquely identifies the hidden state, and the POMDP is block MDP with corresponding Markovian observations.
all w, Pr[w] s;) Pr(ag (r)Pss (w) 0 and thus Pr[w]a,] (()[s)),)r((s(l))))
Next, we consider the case where i = j. We have that for all i, k. Σ Pr(s:|w)) = O. Because each term is nonnegative, this equivalent to Pr(w[s; Pr(a))()) (Pr(s(())
1) = 0. Because we again have that for all observations there exists an action age with nonzero probability, this means we can choose
k/' find Pr(w[a;) = or Pr(s([w) = 1 for all w, st. This means that either the state s; does not produce observation w, the observation w uniquely determines which state the agent is in. For all w and = j. either Pr(w]s;) = 0 or Pr(s,|a) 1, so if we restrict our focus to the set O(s;) {w Ω Pr(w[s;; > 0}. we see that Ewe)(e)) Pr(w[s;) Pr(s)lw)- = EwEo(.)) Pr(w]s;) = For the remaining w & O(s),, Pr(w[s;) = 0.
To recap, the first case tells that for all w and i j. Ever Pr(w]s:) Pr(s)lw)= 0. The second case tells us that for all w and i = j. Ewen Pr(w]s:) Pr(s;[w) 1. Combining both these cases, we see that Σ Pr(w]s;; ) Pr(s)lw) = 1[s; = syl, 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 that the converse is also true. If the system is block MDP, then II$ 4 and so K3 K2¹
19

---

Proof that matching transition-pollcies is equivalent to an MDP.
Assume that TK TKK'. We can reduce this as follows:
0 = TK - TK' T(AII* $ + - >))/W") - T(XII +(1 (A ()I((I
which implies either = X' or TTP TOW!
Claim: if P is POMDP such that 7 TT for any policy, then P is equivalent to MDP.
Recall that the transition function of the effective MDP model is WT and the rewards are W ASa, Intuitively, If and W are equivalent transition functions, it means they induce the same visitation after any number steps. Note that due to the premise, the following quantities are equal:
WISA ToA:
W(TII*) = =
(WT6I):W4,
"M " W
-V
= W7721//ww00000000000000000011111
= =
and so on.
Multiplying each successive equation by 7² and summing the results gives:
W(TSA - = (TRA (
For any given observation, POMDP model and the effective MDP model make the same predic- tions 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 RSA,
W(ISA - -TIII R RSA. (I -W RSA
(V
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 example of such POMDP and its equivalent MDP in Figure
1/4
1/2
1/2
1/2
$12
S4
S2
1/2
84
1/6
1/2
1/2
Figure 7: (Left) A POMDP with six states and actions, where T&W STIS From state 8o- both actions transition to states 81, 82, and 82 with the labeled probabilities. State 8 produces two observations (equivalent those of sj and s2) with equal probability. (Right) An equivalent five- state MDP model that behaves identically to the POMDP for any policy.
20

---

G Memory Augmentation and Norm Details
In this section, we clarify theoretical details in augmenting POMDP with memory and the X- discrepancy norm.
G.1 Memooy-Augmentedd POMDP
As referenced in Section here we will explain how define memory-augmmeteed POMDP from base POMDP (S, A, R, Ωt Φ 7). Given set of memory states M, we will augment the POMDP follows:
SM S SxM M
AM = Ax M
QM ΩX M
TM : SM x AM ASMi
(sy. mu. a1 m1) + T (.: a1) x 8ms (-)
RSuAM; SM x AM + R;
(so. mo. a1 m1) + R(so, a1)
PM 8M AlM:
mu) Φ so) 8m
TM == %
where 8z (-) is the discrete point distribution centered To. and between two distributions denotes the product distribution. To demonstrate what this means, the notation above is equivalent to seeing TM instead
function mapping [0, 1] and defining TM SM AM SSM [0, 1]. (80, mo, mi, 82, m2) + ((82 s0, (1) [mm m2). where a mi.
action
(01, m)) means take action d1 and set the memory
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 AM provides the agent with means of managing its internal memory state. Together, these augmentations allow writing the augmented transition dynamics TM, which are defined so as to preserve the underlying state- transition dynamics T while allowing the agent full control to select its desired next memory state. The observation augmentation 2 provides the agent with additional context with which to make policy decisions, and the observation function 4.M preserves the original behavior of the observation function & while giving the agent complete information about the internal memory state.
We define augmented policy тM and its tensor counterparts IIM and IISM over observations and states, respectively, as follows:
R. Ω. - AAM: 1 is of size (AM x ΩM x Am);
w1 is of size (SM SM x Am):
where
[M Mg, W1, m1, 02, m2) = 1[ww, mo) = (wi, mI (a2, m wi,, m1);
IISM (Sg, Mg. $1, m., 02, m2) 1[(sg. ma) ($), m)) Σ 1 D. m|s0. mg) (a2, m2/ws, m).
(uM
We extend the Pr(s|w) tensors to SM and ΩM as follows:
WM is of size (ΩM x SM))
Will M of size(Si SM Am):
where
WM mo, $1, m.) = 1[mo = m1] Pr(si [wo);
will M (wp, ma, 81, m,, a2 m2) = 1[mm = m. Pr(s (wp) T.M(22 m2/wa, mo)-
21

---

One uppealing aspect of the above augmentation process is that it makes the agent's control over its memory function explicit, via "memory-augman perhaps a bit unusual to combine the agent's s policy and memory together. An equivalent and perhaps more intuitive formulation decomposes this mimtoyyymory policy into separate action policy T S x M - AAA and a memory function Q x M x A + A.::
actions" AM A M.. However, is
(a m' w, m) - T w. w,, m))(mm ww, m, m)
Note that we define memory functions as distributions next memory states for generality.
There are two ways to view this augmentation procedure. In the first view, outlined above, the agent selects both action and a next memory state m', and the environment T'M responds by sampling the next state s' ~ T(1ss, a) while deterministically setting the next memory state to m'. In the second (equivalent) view, the agent selects only action with the agent's memory function behavior folded into the transition dynamics Tμ. This view leads to the following POMDP quantities:
Tp : SM x A ASMi
(so, mo, ao) + T- so. an)us( so, mo, ao);
HSMA SM x A + I;;
(so, mo, ao) - R(50, ao),
where us(-) mo: ao) == Ewer ( (wwo do, mo) is the effective memory function for states.
The p-specific policy and weight tensors are defined as follows:
H. is of size (NM x SM x A); SM is of size (SM x SM x A);
Will is of size(S, SM x A);
where
[wp, mo W1, m1, 02) = 1[(wb, me) =
W1, m1);
I (So, mo: 81, my, 02) = 1[(80, mo) = (81, m1) Σ
$ @M mo)xz((2|u, m);
(w.
Will(wp, (wo, mo: $1, my, 02) 1 1mg m1] (W)) ( ( wwo my). p
To recap, the first augmentation scheme transforms the POMDP - (S, RSA, 4 P.M., into PM (SM AM, 9., TM, MMMA, PM.PMoo 7). The second view forms equivalent POMDP P, = (SM,A, Sm, Tuy MSMA &M. which folds the agent's memory function into the transition dynamics.
Since both of these are valid POMDPs and we already have general expression for the TD(A) value function of POMDP, can immediately write down the corresponding value functions for each of these types of augmentation. For PM. we have:
Qan Qz==
WM WM
TTonAm-TTM(ATSSA + +(1 -1))@mWWH))) -1 :RSSAm;
(9)
and for Pas, we have:
e² = WM (IMA - TTp (AII + (1 w/)))
-RSMA
(10)
We provide pseudocode for taking this memory-Cartesial Cartesian product of a given POMDP in Ap- pendix H.3, Algorithm 3. The pseudocode uses the view that aligns with Equation (10), since that view matches our implementation.
G.2 A- A-Diecrepancy Norm Weighting
The X-discrepancy as introduced in Definition I contains weighted norm ||C ||C+xll, over the observa- tions and actions of the decision process. We use several choices of norm p and weighting C' in our experiments. For p. we use both the p = 2 or T/² norm, and the p = oo max norm, as described below. For C', we have weighte and occupancy weighted cases. In the welihy-weighted case,
22

---

we assign the (w, a) entry of C' the weight (1, x(akw)). In the occupancy weighte case, we assign ((al))), the (w, a) entry of C the weight ((a)))) where Pr(w) proportional to the discounted ob- servation occupancy.
In Figure 4 and Figure 5, use the policy-weighted 1.2 norm. We 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 X- discrepancy puts more weight on the initial observations, which results in memory functions that place more emphasis on resolving x-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 X-discrepancy. as we show in Appendix 1.6 that occupancy weighting in the sample-based setting does not detract from performance in these environments.
In Figure 3, we use the secupancy-weightee max norm 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-weigheer L norm, since it is the simplest choice when using samples as does not require re- re-weighting or importance sampling.
Regardless of which norm we choose, we can estimate gradients as in Equation 7 because:
A A P,
11 2
=
1/2
Σ
Pr(w) Pr(alw)))))) (w,
(4s a))²
w.a
1/2
Pr(w,a) (w, Q (w, a))²
w.a
(B(raa)v
[1)
Q a)) ?)) 1/2
In practice, when optimizing the X-discrepancy, we omit the outer square root and simply compute the mean squared difference between Q-values, since square-root function monotonic transformation and its argument is non-negative.
23

---

H Environments and Experimental Details for Closed-Form Memory Optimization
We now describe the proof-of-concepl experiments from Section 4 involving closed-form memory optimization 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 x 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 X-discrepancy gradients, as well as procedures for policy improvement after learning 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 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, [Iitee) 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,
modified the domain by splitting each state L/R into initial state L,RR, that always emits an initial observation, and a post-listening state Ly/R, that uses the 0.85/0.15 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.
Listen
1.10
Right
Left
Right
L Listen
Listen
100
Left
Left.
100
1/2
-1 Listen
-1 Listen
R
R
RS
Right
Left + +110
Right -100
Lästen
Figure 8: Visualizations of the Tiger POMDP. In the original version (left) the observation function was action-dependent, whereas in modified version (right) observations only depend state. The state color for the domain on the right represents the distinct state-dependent observation functions: purple states the initial observation, while the other states biased towards either left (blue) or right (red) observations with probability 0.85.
H.1 T-maze Details
We use T-maze with corridor length 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." There are 5 observations in this environment one for each of the initial states, corridor observation shared by all corridor states, 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, remains in the current state. From the junction state, the agent receives a reward of +44 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 y = 0.9.
"Technically, we group the four terminal states into single state for conciseness. which is functionally equivalent.

---

+1
Figure 9: Visualizations of value functions computed using MC (left) and TD (right). MC averages over entire trajectories, so can associate the blue observation with the upward goal. By contrast, TD computes value by bootstrapping; its value estimates for subsequent observations ignore any prior history.
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 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 +44 with MC and (4 - 0.1)/2 1.95 with TD.
H.2 Analytical
and Policy Improvement Algorithm
Algorithm describes our memory optimization procedure, which reduces the X-discrepancy to learn memory function, then learns an optimal memory-augmented policy.
Algorithm Memory Optimization with Value Improvement
Input: Randomly initialized policy parameters 0., where IT = softmax randomly initialized memory parameters f,. POMDP parameters P == AS T, RSA Ω. $, Po. 7), number of memory improvement steps pu number of policy iteration steps nsteps, learning rate E [o, 1]. and number of initial random policies n.
Output: Optimized eeemory-augmente policy parameters On, and memory parameters 0,.
Il Initialize random policies and select policy to fix for memory learning.
+ 1 handomly.iittm..ooiiiiee n.poolcces(m)) 0, select argmax lambda discrepancy (0....,, bn-1))
0,-1)
Repeat policy over all memory states.
0.. repeat(O,, [M])
Improve memory function.
memory -eemoryyimproveeeeeeeeeeeeeeeeee _improvemen((f... Oz,, , P, nateps..)
Augment POMDP with memory function.
Pp + expand_over wet.ooerrmmmmmyee 0g)
Improve memory-augmenteed policy over memory-augmentedd POMDP
pxp t-
(f (0,,, Pps neseps.:))
return 0g ON
In Algorithm the rameterized policy. We use an analytical version of the policy gradient [Sutton et al., 1999] algo- rithm. The wtttoooly-iiittnoollll select.
function can be any function which improves a pa-
function returns randomly initialized policies, and the lambda discrepancy function picks the one with the largest X-discrepancy. The
argmax select.argmax memory hemory-improveneee improvement and expand_oves over memory functions are defined in Appendices H.2.1 and H.3, respectively.
We have noticed that larger X-discrepancy tends to provide better signal for learning memory functions. Sampling random policies for memory improvement highly likely reveal a X- 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 X-discrepancy as the basis for memory optimization.
25

---

H.2.1 Memory Improvement Algorithm
In this section we provide pseudocode (Algorithm 2) for the memory-learning algorithm described in Section and used in Algorithm 1. This function takes input POMDP P and memory parameters 0,. and minimizes the X-discrepancy of Definition 1. This minimization is achieved through gradient descent update computed using the auto-differentiation package JAX [Bradbury et al., 2018].
Algorithm 2
Memory Improvement
Fixed policy parameters OZ., where 1 SM x A) is the expanded effective policy over state for the policy Fn = soft sax (0=,) (SM x A), memory parameters 0,., POMDP parameters P = (S,A, BSA,, Ω2 2 Po: 7), number of improvement steps nstess learning rate [0, 1].
Input:
Output: Optimized memory parameters Ap-
for i = O to Ts do // Augment POMDP with memory parameters 0y (SM,,, A, Tus Ry. RSMA, &MM P.Mo: & expand nemory(P,0.. II Compute TD value function (with memory augmentation)
Q4
np = =WM -T, MM W
II Compute MC value function (with memory augmentation)
Q+ np = =WM -I-TTMII55 M
I1 Calculate the X-discrepancy
A y =1[0a, 2t 1|,22
II Calculate the gradient of Ar, w.r.t. pn update memory parameters
02 <- update.paramse 0p Vo, A a.)
end for return 0
Here, update-params () is any seefient-deeeent-lee update, such stochastic gradient descent, Adam, etc. As a note, all parameters in these experiments are initialized with a Gaussian distribution, with mean and standard deviation 0.5.
H.3 Memory Cartesian Product
In this section, define the memory- memory-Cartesian Cartesian product function, expand _over www.orymmoc..... () used by Algorithms and 2. This function computes the Cartesian product of the POMDP P and the memory state space M, as described Appendix G.I.
Algorithm 3 Memory Cartesian Product (expand over memory)
Input: Memory parameters 0, (with corresponding memory function μ1) POMDP parameters
P = (T, RSA, 4 po.?), 7), number of memory states [M]
Repeat reward function for each state over each memory m E M (SM A).
R RSMA t- repeat over states (RSA, (M))
Calculate transition function memory cross product.
11 First, calculate the effective memory function over state (S x A M M).
us & einsum ij, jklmm - iklm', p,1)
I1 Now expand the state transition function to include memory state transitions (SM x A x SM).
Ta & einsum( cinsum('iljk, jij lim - lijmn', ps,
Il Calculate observation finntttoooooonnii memory cross product (SM x P.m). IMMi is the identity function over IM].
$ MM + -krooon ($) M))
Finally, calculate the initial state distribution (SM)-
PMo & (po(s) (s) if m = O else for s, m E S.
return Pμ (SM. A. S.. T,, RSMA, $ PMM.PMo.
26

---

Note that einsum is the Einstein summation, and kron is the Kronecker product. The augmented initial state distribution is simply the same distribution as Pp- except with probability mass over all non-zero memory states, since the memory state always initializes to memory state O)
H.4 Closed-Form Memory Learning Experiment Details
For all experiments in Section we run memory optimization on the suite of POMDPs with the following hyperparameters. We optimize memory for nseess, 20K steps and run policy iteration for 7steps, 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 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 metric comparable to the initial state distribution weighted value function norm, which we use as a performance metric for our memory-aaugmented agents.
The belief-state solution for the x 3 maze was solved using an epsilon parameter of 0.01, due to convergence issues with the environment when utilizing the POMDP solver.
H.S Closed-Form Experiments on Parity Check
In Section 3.3 we introduced the Parity Check environment as an example of POMDP with zero x-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 be quite rare. The symmetry disappears if modify the example even slightly, such as by changing the start state probabilities or by introducing "stay-in-place action at any one color. These modifications lead the x-discrepancies Figure 10, which offer some reassurance that such edge cases are the exception and not the rule.
Besides modifications the POMDP itself, adding memory induce non-zero X-discrepancy. In particular, random memory functions reveal X-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 (right) followed almost the procedure as Algorithm 1 except for selecting the initial policy. In this case, any randomly initialized memoryless policy has zero X-discrepancy, but augmenting with a randomly initialized memory function reveals non-zero measure. Thus, for this experiment, augment our POMDP with a random memory function before choosing the initial policy that maximizes the X- discrepancy. All other algorithmic details remained the same.
101
109
>-discrepancy
A-discrepancy
10-6
10-6
10
10-12
(r(STAY|BLUE)
Po PRED - PINK)
10
18
10-18 18
0.0
0.5
1.0
0.0
0.5 Probability
1.0
Probability
Figure 10: (Left) The Parity Check environment (reproduced from Figure 4). (Right) Minor mod- ifications to the transition dynamics (left) or initial state distribution (right) result in non-zero discrepancy for almost all policies.
27

---

I Scaling the X-Discrepancy with PPO and RNNs
L.1 Algorithm Details
We build our X-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.
(0) = E, (OActor 0RNN) cv Lv (Peritics RRNN) + CEnt / [T((- Wr, 2)]]
(11)
Here, LOLP is the clipped surrogate loss for the actor, defined by:
LCLIP BRNN) Ex min(p: -0tp://0)))_..cppppp0011--------
(12)
where Pe(0) is the ratio between policy probabilities of the old policy and new policy, A, is the estimated advantage at timestep f, and cv and Lv are the value coefficient and value function loss, respectively. Lv is the mean-squared TD error, with targets calculated using truncated version of X- return [Sutton and Barto, 2018]. Finally, CEnt and [T Wr, -t-1)] are the entropy coefficient and entropy of the policy Ty 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 non- recurrent architecture that is function of only the current observation (and optionally the most recent action).
Our proposed algorithm replaces by with A2 which has three components: two value losses for two separate value heads (each estimating X-returns for different As), and our X-discrepancy auxiliary loss.
LV, d2 (fCritic: BRNN) = (0) + (1-3) (fCritiet, RRNN) + Lv, (foriiies BRNN)),
where LA (0) is our X-discrepancy loss, defined in Equation 8, and Lv. is the A-return-asttrree as-ret TD error for a given A. Here, 3 E [o, 1] is a hyperparameter trading off between the X-discrepancy loss and the value losses. Finally, the agent only uses Vx, its value estimate for its advantage function.
The computational and memory requirements for PPO with and without the X-discrepancy auxiliary loss similar. While the X-discrepancy loss requires training 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 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.
I,2
Environment Details
All four of the environments used for evaluation in Section were all ie-mmleemettaionss of envi- ronments used in Silver and Veness [2010], in JAX [Bradbury et al., 201 8), allowing for massive hardware acceleration. We now give details of these environments.
L.2.1 Battleship
Partially observable Battleship, one-player, limited observation variant of the popular two-player board game, was first introduced as 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 at every step, the agent receives a binary observation: 0 if the last shot missed 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, (1..... 10} x {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 horizon at most 10 x 10- = 100. The environment terminates when all positions on the grid with a ship are hit. The rewards at every step are -1, with positive reward of 10 x 10 100 when all ships are hit/the environment terminates. The discount factor y is set to 1 here, since we are in the finite horizon setting.
= {0, 1):
At every environment reset, 4 ships of length (5, 4,3, 2) are uniformly randomly placed on 10 x 10 grid. At every step, to allow for easier learning of all agents, we concatenate the last action to the observation, that our RNN conditions on the entire history, opposed to only observations.
28

---

1.2.2 Partially Observable PacMan
Partially observable PacMan (a.k.a. PocMan), 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, 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 1 11-dimensional binary vector. The first four values are 1 if there a wall in each of the four cardinal directions of the agent. The next element is if there is food Manhattan distance 2 or less away from the agent. The next value is if there is ghost Manhattan distance or less away. The next four elements are 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 x 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 y = 0.95.
The episode terminates when the agent is killed by 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 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 top of the PacMan environment in the Jumanji reinforcement learning framework [Bonnet et al., 2024].
I.2.3 RockSample
RockSample [Smith and Simmons, 2004], a well-known partially observable benchmark in POMDP planning literature, is rock-collecting problem simulating a rover scanning potential rock samples and collecting them if they are desirable. RockSample (n, is a problem with a grid of size n x n, and k rocks distributed randomly throughout the environment. The two variants we test have grid size of 11 11 and 15 x 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 + + 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 ie rocks. Checking rock will probabilistically tell the agent the correct parity of the rock, depending on the half-efficiency distance dhed and the l2 distance ds(s) between the agent's position and the rock being checked, given the current state s:
Pr (accurate 8, a check;) 1 (1 21+2-d(t))dann
(13)
In this work, we simply set dhed to be the maximum possible distance between any two points in the grid. Traversing to and then sampling a good rock gives reward of + +10. Sampling bad rock gives a reward of -10. Exiting to the east border of the environment will terminate the environment and result in reward of + +10. After sampling a good rock, sampling the same rock again will result in reward of -10. This environment has a discount rate of y = 0.99
The 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 vector of size 2n + binary values, where the first 2n elements are a two-hot representation of the zyycoordinates of the agent. The final i observations are set after choosing a check action, and set 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 (TTo et al., 2023], as opposed to model-based planning algorithms. The first is that after an agent checks a rock and the agent receives positive sensor reading, the observation bit corresponding
the position of the
29

---

to this rock is set to 1, and remains I 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 rock, the observation bit corresponding the sensor reading of the rock remains O. This makes the problem harder, since the agent has to infer the negative sensor reading, and disambiguate it from null sensor reading, when the agent takes non-check actions.
30

---

1.3 Discounted Return RNN Results
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 finite-horizon and undiscounted.
RockSample 11 x 11
RockSample 15 x 15
Discounted Return
20
10
0.0
0.4 Timestep
0.8
Timestep
x106
x 107
Battleship 10 x 10
Partially Observable PacMan
Discounted Return
300
15
200
PPO RNN LLD
10
PPO+RNN
100
PPO
0.0
0.2
U.4
0.6
U.S
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Timestep
x10
Timestep
x x10"
Figure 11: The x-discrepancy (LD) performance in discounted returns over recurrent (RNN) and
memoryless PPO. Learning curves shown are the mean and 95% confidence interval over 30 runs.
RockSample 11 11
RockSample 15 15
Episode Reward
20
20
0-
2
3
0.0
0.2
0.4
0.6
0.8
1.0
Timestep
x x10
Timestep
x10
Battleship 10 10
Partially Observable PacMan
Episode Reward
15
2000
RNN
10-
1000
PPN+
PPO
5-
0.0
0.2
0.4
0.6
0.8
1.0 x 107
0.0
0.2
0.4
0.6
0.8
1.0
Timestep
Timestep
x x10
Figure 12: Larger version of Figure 6. The X-discrepancy (LD) performance (undiscounted returns) over recurrent (RNN) and memoryless PPO. Learning curves shown the mean and 95% confidence interval over 30 runs.
31

---

1.4 Experimental and Hyperparameter Details
Our base PPO algorithm 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 al., 2022] with hardware acceleration. The hyperparameter sweep was performed on cluster of NVIDIA 3090 GPUs, and the best seeds presented were run one GPU for each algorithm, running for I to 12 hours, depending on the domain. Our environment vectorized over copies, and we use truncated-backe bacpeetootcckkooptttoooo through time [Jaeger, 2002] as our gradient descent algorithm, with truncation length of 128. Our LCLIP clipping is set to 0.2. The value loss coefficient is set to cv 0.5. We also anneal our learning rate over all training steps, and clip gradients when the norm larger than 0.5 by their global norm [Pascanu et al., 2013].
We now detail the hyperparameters swept across all enviromments and all algorithms. We do so in Table
Hyperparameter
Step size
[2.5 1053, 10 1,-4,25 0.7, 0.1, 0.7, 0.9, 0.95] B (A-discrepancy)
10 5,2.5 10 10
A
X2 (X-discrepancy)
00, 0.155, 0.25, 0.5]
Table 1: Hyperparameters swept across all algorithms. Rows labelled with X-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 b-diserepancy-augmeee 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
recurrent
Step size
À1
Xa
B
Battleship
2.5 10
0.1
0.95
0.5
PocMan
2.5 10
0.95
0.5
0.5
RockSample (11, 11)
2.5 10-4
0.1
0.95
0.5
RockSample (15, 15)
2.5 x 10-4
0.1
0.5
0.25
(a) Best hyperparameters found for cesssrrppanc--augmenne recurrent PPO.
Step size
A1
Step size
\1
Battleship
2.5 x 10-5
0.7
Battleship
2.5 x 10
0.1
PocMan
2.5 10-5
0.7
PocMan
2.5 10
0.7
RockSample (11, 11)
2.5 10-4
0.7
RockSample (11, 11)
2.5 10-4
0.1
RockSample (15, 15)
2.5 x 10-5
0.5
RockSample (15, 15)
2.5 1
0.1
(b) Best hyperparameters found for the recurrent PPO basclinc.
(c) Best hyperparameters found for the memoryless
PPO bascline.
Table 2: Best hyperparameters for each environment and each algorithm. Hyperparameters were found using seeds, and taking the maximum AUC.
Our network architectures are standard multi 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 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 2x latent size for this environment), that takes input the outputs of the previous layer
32

---

concatenated with the hit-or-miss bit (a residual connection). For our memoryless baseline, we replace the recurrent neural network with 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. uses this new latent state as inputs to both the actor and critic.
We now detail environment-speciln hyperparameters in Table 3. All hidden layer latent sizes are the same sizes as the latent sizes. The full implementation of algorithms, experiments, and environments are available at https: : / //github. com/br htt/bbowwirll./m www.wrmlllmbbaaaaeeeeeee _ammbdddicreeanccy
Latent size
CEnt
Battleship
512
0.05
PoeMan
512
0.05
RockSample (11, 11)
128
0.35
RockSample (15, 15)
256
0.35
Table 3: Environment-<pecifie hyperparameters, across all algorithms. We set the entropy coefficient to a higher value in RockSample because the environment requires more exploration.
1,5
P.O. PacMan Pellet Probe Visualization Details
We train two RNN hidden state probes in order to generate the memory visualizations in Figure 2. Probes were trained on the hidden state outputs of the RNN PPO and RNN PPO LD agents. 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 collect both RNN + PPO hidden states and the RNN + PPO + LD hidden states. With this dataset, the each probe trained with the corresponding RNN hidden states 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 sigmoid function over the final logits to map outputs to 0 and 1. We use binary cross entropy loss between these predictions and the pellet occupancy targets. The hidden size of the network was 1024, with a step size 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 training reach the performance visualized.
166 Small-Scale POMDP Experiments
We also run both recurrent PPO and our X-discrepancy- eaaisreepancy-aggmmeteedd PPO algorithm on the small-scale POMDPs evaluated in the analytical experiments in Section 4. For these experiments, all agents had latent state size of 32, with CEnt 0.05, and the same hyperparameters swept 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, with the results in Figure 6, adding our X-discrepancy auxiliary loss never harms the performance of PPO.
33

---

4x3
Cheese
Hallway
0.5
0.50
0.0
0.50
0.25
-0.5
0.25
0.00
0.0
0.5
0.0
X10
x
X10
Heavenhell
Network
Paint
Discounted Return
1.0
0.5
0.5
0.0
0.0-
-000
0.0
X10
X100
XIF
Shuttle
Tiger
T-maze
20
-20
PPO+RNN+LD
PPO+RNN
0-
0.0
X10
X10
XI
Timestep
Figure 14: The X-discrepancy performance over recurrent PPO on small-scale POMDPs. Learning curves show mean and 95% confidence interval over 30 seeds.
34

---

NeurIPS Paper Checklist
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 "the discrepancy is exactly zero for all Markov decision processes and almost always non-zero for 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- results Section 5 show that minimizing the X-discrepancy between two learned value functions an auxiliary loss alongside traditional reinforcement learning, is often significantly more effective (and never worse) than training the same agent with just 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 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 other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
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 encouraged to create separate "LLimitations" section in their paper.
The paper should point out any strong assumptions and how robust the results 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 the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or 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, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
35

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Theorem 1 is briefly justified just below the statement and is given in full in Appendix D. We prove Theorem in Appendix E. We also prove one additional claim in Appendix F. We explain the memory augmentation process in Section G.I.
Guidelines:
The answer NA means that the paper does include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question:
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data provided not)?
Answer: [Yes]
Justification: All details of experimental and environment set up are described in depth in the paper. to allow for reproducibility of all experiments, Appendix H gives full detail of the proof-of-concept experiments run in Section 4, and all details of the experiments in Section are described in Appendix 1.2 (for the environments) and Appendix I.4 (for the experiment set up and details). Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
36

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome 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.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: As stated in Section 1, code is available https://gii thub...tt.u....m com/ brownirl/lambda. httaaaaabaddddpppaaeaeaa discrepancy. This repository both has detailed instructions on how to conduct the sweeps, as well the best runs over the 30 seeds mentioned 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.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
Seaiddes/Codeuubmiss
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/
htteeee//oooSeeeeeeeeeeeee asionPolicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperpa- rameters, how they were chosen, type optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: These are present in Appendix H and Appendix Specifically, in L.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.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
37

---

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 enror bars, confi- dence 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).
should be clear whether the error bar is the standard deviation 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 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 range (c.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.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details of compute usage, hardware and time taken for experiments are discussed Appendix 1.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 compute required for each of the individual experimental runs as well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https://neurips.. httpllpp..cpppppppppppep httpppulllctiiicsiiieeee..
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, 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 com- pared 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 Section and Section 5. There is sensitive data requiring strong security measures. Standard
38

---

security measures are in place 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: //gi thub. 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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e-g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: We discuss the broader impacts of this paper in Appendix B. Our paper, as a work the foundations of reinforcement learning, is not tied 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), faireess considerations (e.g., dep.lymenn 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 tied to particular applications, let alone deployments. However, if there is 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, 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work focuses 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 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.
39

---

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Open-source repositories used in this work have been cited, particularly in Appendix 1.4. In addition, our code open- open-sourced with the Apache 2.0 license.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URI.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The released codebase well documented, with proper READMES on experi- mentation and data parsing.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Neither crowdsourcing nor research with human subjects was required.
Guidelines:
40

---

The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation,
or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Neither crowdsourcing nor research with human subjects was required.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere 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.
41


Reviews:
[{"review_id": "UVYgthHQrw", "reviewer": "Reviewer_pbjp", "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.\n\nThe 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.\n\nFurther, 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.", "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?\n    \n- 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?\n    \n- Q3: How is the LD method affected when switching between online and offline-RL settings?", "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.\n    \n- 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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The paper bridges nicely from theory to empirical experiments. The readers are first taught the theory and origin of the Lambda Discrepancy. Then the paper explains the possibility to use LD to detect partial observability before extending this possibility by integrating the discrepancy into the loss function of a policy optimizing algorithm. -> Great story telling.\n    \n- The LD method is versatile and can be used for any optimization algorithm using value estimation. Hence, due to its effectiveness and limited added complexity to an optimization algorithm, the Lambda Discrepancy is an effective way to improve existing optimization algorithms.\n    \n- The experimental results are used in a concise and convincing manner. The language is clear and to the point.", "weaknesses": "- 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.\n    \n- 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.\n\n- The “Limitations” section should be part of the main corpus.\n    \n- 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.\n    \n- 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.\n    \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your response and feedback! We are glad to hear you thought the writing was “clear and to the point.”\n\nWe will respond to your comments in order.\n\n---\n\n> *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.*\n\nThank you for the suggestion. While we use RNNs (specifically GRUs) in our experiments due to their popularity in RL, there are other architectures we have not considered. Transformer-based architectures have recently been studied in reinforcement learning settings, and have been shown to help in long-term memory tasks [1]. Seeing as our metric poses no constraints on function approximation or network architecture, an interesting direction for future work would be to see if the LD objective would benefit Transformer-based architectures as well.\n\nAs for methods that explicitly consider partial observability with memory, the most dominant approaches have used belief-states. Belief-state methods [2] are a class of solution methods for POMDPs that have seen substantial success in robotics [3], but most approaches are intractable for even small environments where the dynamics are unknown, since most methods rely on some form of state information to form the belief state that acts as memory.\n\n- [1] Ni, Ma, Eysenbach, Bacon. NeurIPS, 2023.\n- [2] Kaelbling, Littman, Cassandra. Artificial Intelligence, 1998.\n- [3] Thrun, Burgard, Fox. Probabilistic Robotics. MIT Press, 2005.\n\n---\n\n> *Ideally, it would be most interesting seeing results for environments with Markovian state representation that gradually shift to POMDPs that are not Markovian.*\n\nThanks for the suggestion. We ran this experiment for T-Maze and included the results in our global response.\n\n---\n\n> *The “Limitations” section should be part of the main corpus.*\n\nWe had originally moved the limitations section to the appendix due to space constraints, but we will move this section back into the main corpus with the extra page allotted for the camera-ready.\n\n---\n\n> *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.*\n\nWe proved block MDPs have LD=0, and this is good, because block MDPs are fully observable; they are MDPs. In fact, classic MDPs can be viewed as block MDPs with a single observation per state. Block MDPs are therefore not an edge case, but an indication that our metric is correctly labeling MDPs as having zero partial observability.\n\n---\n\n> *However the explanation of “Intermediate values of λ smoothly interpolate between the two” (pdf) (l. 163) values is not intuitive. An example at would be helpful*\n\nWe apologize for the confusion. You are right that TD(λ=0) and TD(λ=1) don’t provide much intuition on how the intermediate lambda values will behave. For intermediate λ values, TD(λ) is a weighted average of all the n-step returns. The λ parameter controls the weighting scheme: lower values of λ give more weight to the shorter rollouts. When λ=0.5, this defines a weighting scheme where each n-step rollout contributes half as much to the weighted sum as the (n-1)-step rollout. It is possible to view TD(λ) as operating in either a backwards-looking or forwards-looking way, but in the paper, we always take the forward view: the λ parameter does not determine whether previous states are taken into account, but rather how much each n-step future return contributes to the weighted average.\n\n---\n\n> *The experiments remain on a low complexity levels*\n\nThanks for the feedback. Please see “Environment Complexity\" in our global response above.\n\n---\n\n### Questions\n1. *Could you please clarify the setup regarding the LD loss?*\n\nWe have provided a visualization of the network architecture in new Fig. 2 (left). We train two value heads and a policy head on top of a single shared recurrent representation. Only one of the value heads is needed for the PPO objective, but both are used in computing the LD loss, and the recurrent representation is trained end-to-end. For more details, see Appendix I. We will add the figure to the camera-ready if the paper is accepted.\n\n---\n\n2. *How does the LD method benefit when there is existing domain knowledge about an environment?*\n\nGreat question! Domain knowledge can help a system designer to provide useful (ideally Markovian) features to the decision-making agent. If the system designer wants to check whether these features are in fact Markovian, he or she could use the LD to test for it.\n\nInterestingly, even when the environment is fully observable, the λ-discrepancy may still be useful, because function approximation can sometimes act like partial observability. If the environment has Markovian observations but these observations are sent through a lossy function approximator (like a neural network), then minimizing λ-discrepancy could help the function approximator to learn better features.\n\n---\n\n3. *How is the LD method affected when switching between online and offline-RL settings?*\n\nAnother good question! For offline RL, LD needs to be applied with care when the data is not generated from the policy being evaluated, TD(λ) no longer has the same fixed point for all values of λ under different policies. Our theory relating the Markov assumption to LD would not hold between two value functions under different policies. However, one could use tools from off-policy RL such as importance sampling to adapt our methods to this new setting. Another idea is to minimize the LD for the data-generating policy instead of the evaluation policy. Due to the result in Theorem 1, this should result in a memory function that, in theory, mitigates partial observability for other policies as well. This is similar to what we do in the analytical experiments in Sec. 4. That setup is akin to an offline policy evaluation approach for memory learning, which is an exciting avenue for future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "k4wkPQtp0p", "reviewer": "Reviewer_rpiT", "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.", "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?\n2. 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.\n3. 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?\n4. 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?\n5. 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?\n6. How does the memory learnt by this approach compare with representations learnt by POMDP based approaches such as PSRs etc.?\n7. 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": "The authors address the limitations of their work in a separate section in the Appendix.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and the claims and explanations are clearly given. The new $\\lambda$-discrepancy metric proposed by the authors is an interesting one and the entire approach is novel and useful.", "weaknesses": "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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and feedback. It seems that your main concern is the misconception that LD “requires operating over the entire state, observation and action space for its computation.” We apologize if this was confusing. While the *derivation* of the metric uses tensors that operate over the entire state, action and observation space, the actual estimation and minimization of the metric only uses estimates of TD($λ_1$) and TD($λ_2$), which are computed directly from sampled (observation, action) pairs. This means that our method does not require the underlying model of the POMDP to mitigate partial observability, only samples.\n\nWe realize this may not have been clear in our submission, so we’ve added the following discussion near the end of Sec. 4:\n\n> Despite its definition in terms of $S$, we can express the $λ$-discrepancy in a form amenable to estimation from observable quantities only. By selecting the appropriate weighted norm, the $λ$-discrepancy becomes an expectation over observation-action pairs (or, in the case of memory, observation-memory-action tuples):\n$$Λ\\_{\\mathcal{P},π}^{λ\\_1,λ\\_2} := \\|\\|Q\\_π^{λ\\_1} - Q\\_π^{λ\\_2}\\|\\|\\_π = 𝔼\\_{π}\\left[ Q\\_π^{λ\\_1}(ω, a) - Q\\_π^{λ\\_2}(ω, a) \\right]$$\nThe 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.\n\nWe hope by clarifying our objective's dependence on only observables we've highlighted the broad applicability of this method in partially observable domains.\n\n---\n\nWe will now respond to your questions.\n\n1. *However, the authors also use two RNNs in their proposed approach that uses the entire history. So how does this argument hold?*\n\n    Sorry for the confusion. We meant to say “Without state, a model-based comparison would require variable-length history inputs **per time step** to ensure a Markovian representation”. So the input size for this model-based comparison would be on the order of O(t), whereas the input size to an RNN is fixed and always O(1). To clarify, we only use a single RNN, with two value *heads* conditioned on the same internal representation (new Fig 2, left).\n\n2. *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.*\n\n    Good question! We address this with a new experiment in the global response (“Parity Check”).\n\n3. *How will the convergence analysis of PPO change after the addition of the auxiliary loss given by Eq. (8)? When these functions are being learnt (still have errors) the λ-discrepancy will also have an error. How do the authors address this?*\n\n    While we are unaware of any general convergence guarantees for PPO (much less recurrent PPO), our empirical results show that the LD loss frequently performs better (and never performs worse) than the RNN + PPO baseline.\n\n    You are correct that the value functions will be inaccurate during learning, and so too will the LD. Nevertheless, just as value functions are commonly useful for RL even before they have fully converged, so too is the LD, as our experiments demonstrate.\n\n4. *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?*\n\n    Our approach is introducing a new metric, the LD, as opposed to advocating for a particular algorithm like PPO. The LD is not limited to only PPO, but could be leveraged in essentially any RL method that uses value functions. We are considering extensions to other approaches for future work, where we believe the LD still has great potential to improve existing algorithms.\n\n    We aren’t sure what you mean by state-of-the-art POMDP methods. If you mean belief-state methods like POMCP, it is unclear what such a comparison would reveal, since those methods assume knowledge of the transition, reward, and observation functions, and the set of hidden states. Our method only assumes access to on-policy observations while acting in the environment, which is a much more challenging setting.\n\n5. *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?*\n\n    The proposed algorithm requires two value function *heads* (3 layer MLPs), each conditioned on the hidden state of a single RNN. The extra computational overhead is only for this additional value head---about a 12% increase in parameters vs. the RNN baseline, as we discuss in Appendix I. We feel this is a reasonable tradeoff, given the performance gains.\n\n6. *How does the memory learnt by this approach compare with representations learnt by POMDP based approaches such as PSRs etc.?*\n\n    Good question! We added a new experiment in our global response visualizing these memory functions for the methods in Sec. 5.\n\n    Our approach uses memory functions to incrementally update a summary of past observations. By contrast, PSRs use tests of potential future trajectories, where the main challenge lies in finding these tests. In prior work [Singh, Littman et al., 2003], PSRs were able to learn representations for domains as large as 58 states and 21 observations. By contrast, Battleship (without accounting for possible ship locations) has ~2^100 states. PacMan, even when ignoring food dot configurations, has ~5^400 states. To the best of our knowledge, PSRs have not yet been scaled up to domains of the size we study here.\n\n7. *Can λ-discrepancy metric be defined at a state/observation or state/observation x action level? If not, it is not clear how the quantity given in Eq. (8) relates to the one in Definition 1?*\n\n    We hope our initial clarification above answered this question!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ckmMLGzHRA", "reviewer": "Reviewer_AeCz", "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.", "questions": "- What does LD stand for in Figure 4?\n- 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?\n- 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": "The authors have adequately addressed the limitations of the paper.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The idea is novel and reasonable. There is a natural connection between $\\lambda$-return and the Markov property of the underlying environment. By minimizing the $\\lambda$-discrepency using RNN policies, the learner is able to extrapolate a state representation that can sufficiently summarize the past history.\n- The concept of $\\lambda$-discrepency is novel and could be useful for RL researchers working on other policy learning problems.\n- The authors demonstrate the performance improvement on a select set of complex POMDP environments. This suggests that the proposed approach could impact some practical domains.", "weaknesses": "- 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?\n- 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.\n- Additional discussion on why minimizing $\\lambda$-discrepency could improve the planning performance would be appreciated. Could the authors provide some intuitions?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback and questions.\n\nWe will group two of your comments together, and respond to the remainder in order.\n\n> *As for Theorems 1 and 2, do the policies π have to be Markov policies only taking the current state S as input? If so, the underlying environment is expected to be a block MDP. I am wondering if the result could be more general if we allow non-Markov policies with memory.*\n\n> *Q3: For policies π in Theorems 1 and 2, are they Markov policies only taking the current state $S$ as input? If so, it excludes the non-Markov policy, which could have memories.*\n\nThere seems to be a minor misconception. The agent’s policy $π(a|ω)$ is conditioned on observations---not hidden state. Rather, the *state policy tensor*, $Π^S\\_{s,s',a}=𝟙[s=s']\\sum\\_ω \\Phi\\_{s,ω}π\\_{ω,a}$, is the *effective* policy over states; this tensor maps states to the expected policy under their observation distribution. At no point do we require the agent’s behavior to be conditioned on unobservable quantities. However, our definition does also support policies that condition on both the observation and the memory state, via the memory augmentation process we outline in Sec. 4.\n\n---\n\n> *Could the authors make the connection with the weakly revealing POMDPs?*\n\nWeakly revealing POMDPs (Liu, Chung, et al. 2022) are an interesting class of POMDPs in which an $m$-step future contains enough information to reveal the current state, given sufficient samples, and the paper authors argue that this condition makes partial observability “not scary”. At a high level, the LD also helps to make partial observability “not scary”, by providing a way to detect and mitigate it. So it would make a lot of sense if the weakly revealing literature ultimately turns out to have some connection to this work. We have not been able to prove a formal connection yet, but if you have suggestions on how to do this, please let us know! Ultimately, we hope that the LD will be a useful tool for analyzing many aspects of partially observable reinforcement learning.\n\n---\n\n> *Additional discussion on why minimizing λ-discrepency could improve the planning performance would be appreciated. Could the authors provide some intuitions?*\n\nCertainly! Lower LD leads to state representations that are more predictive of the future. Better representations lead to better value function estimation, and this leads to better performance. Why does it help to minimize this LD quantity specifically? Well, to the extent that optimal behavior requires knowledge of the underlying hidden state, zeroing the LD is a reliable way to ensure that the agent’s representation contains a complete picture of that hidden state. On the other hand, if optimal behavior only requires conditioning on the current observation, then LD minimization won’t improve performance, but it can still provide additional confidence that the current observation is indeed sufficient.\n\nTo see what these “better” representations look like, check out “Memory visualization” in our global response above. We also find that LD has a minimum at zero and increases with partial observability (see ‘Varying “Markovness”’). We also conjecture that reducing LD also improves value estimation error. We have not yet proved a formal link, but this seems plausible, particularly given that LD is defined as the difference between Q-values, which is essentially a type of value error. Moreover, the empirical results in our paper suggest we are indeed learning better representations, which should allow better value estimation.\n\n---\n\n### Questions\n\n> *Q1: What does LD stand for in Fig. 4?*\n\nSorry for the confusion. LD stands for the Lambda Discrepancy in this Fig. 4. We’ve updated the caption to clarify this.\n\n---\n\n> *Q2: How to select parameters λ_1 and λ_2 when minimizing the loss function in Eq. 8? Does one have to try multiple pairs of such λ s?*\n\nThanks for the question. Please see “Understanding and selecting λ parameters” in our global response.\n\n---\n\n> *Q3.*\n\n(Addressed above.)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gpm23z1GDS", "reviewer": "Reviewer_8yd3", "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.", "questions": "Please see weaknesses above.", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "**Originality**: I commend the authors on their vision described in this paper. The lambda discrepancy seems to be widely applicable, scalable, and addresses an important problem broadly relevant to sequence decision making. It leverages well understood techniques in the field (namely, TD-lambda) in a creative new way. To my knowledge, the idea is novel and unique. \n\n**Quality**: The use of the lambda discrepancy as a measure of partial observability is theoretically justified. It is shown that the measure is equal to 0 under a Markovian state representation, while typically being non-zero otherwise. The experiments appear sound and show optimizing the lambda discrepancy aids in learning a good policy.\n\n**Clarity**: See comments below.\n\n**Significance**: I believe this idea has the potential for very high impact. It is general and targets a challenging and important problem.", "weaknesses": "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). \n\n**Quality**: The discussion of some important points -- both theory-wise and experiment-wise -- is missing. In (rougly) chronological order:\n\n(Theoretical:) \n- 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.)\n\n- 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.\n\n(Experimental:)\n- 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].\n\n- 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.\n\n**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:\n\n- 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$. \n- 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$.\n- 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? \n- 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. \n\n**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. \n\n[1] Zhang, Amy, et al. \"Learning Invariant Representations for Reinforcement Learning without Reconstruction.\" International Conference on Learning Representations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the kind words and feedback. We're glad you had such a positive impression of the idea and see potential for very high impact. You cited clarity as your biggest concern, so we'll address that here and put our other comments in the global response.\n\n**Markovian Representation**\n\nWe apologize for the confusion. There was a typo in line 97. It should instead read \"next representation z_{t+1}\". The definition can be applied to states (s), observations (ω), and even memory-augmented observations (ω, m) by replacing z with the relevant quantity. If equation 1 holds, the corresponding representation is Markovian.\n\n**Empirical MDP**\n\nYou're correct that Pr(s|ω) is an expectation over trajectories under the policy. Assuming all observations have non-zero probability, we can write this in terms of trajectories $\\tau=(s_0,ω_0,a_0,r_0, s_1,ω_1,a_1,r_1,...)$:\n\n$$Pr(s|ω)=\\frac{𝔼\\_{\\tau\\sim π}[\\sum_{t=0}^∞ γ^t𝟙[s_t=s,ω_t=ω]]}{𝔼\\_{\\tau\\sim π}[\\sum_{t=0}^∞ γ^t 𝟙[ω_t=ω]]}.$$\n\nThis is equivalent to the definition we use in Appendix C, but we agree it is easier to understand it this way up front.\n\nYour point that the empirical MDP is “not a real MDP” is interesting. MDPs are simply models of environments and can be of arbitrary quality. The empirical MDP is one such model. It is the one practitioners use implicitly whenever they treat POMDP observations as Markovian \"states\". The empirical MDP model consists of (Markovian) transition and reward functions that condition on observations, which we define in terms of Pr(s|ω), and which, crucially, might make different predictions from the POMDP model. We can thus view Pr(s|ω) as a weighting scheme for aggregating the behavior of the states that produce each observation.\n\nModel selection is not a new idea, but our specific characterization of the empirical MDP for a given POMDP is novel to the best of our knowledge.\n\n**Tensor Notation and Q vs. V**\n\nDefining LD using probability notation directly would be prohibitively cumbersome and a poor use of space (see Appendix C). Tensor notation succinctly expresses the concept and has precedent in many papers [Lagoudakis & Parr 2003, Tagorti & Scherrer 2015, Gehring et al. 2016]. Q-values are easy to express with tensors that track action and observation indices separately, and they are less likely than V-values to succumb to Parity Check symmetry, since they avoid an additional summation over actions.\n\n### LD=0 cases and Parity Check\n\nWe found a new way to characterize these edge cases that streamlines the section, removes the need for Eq. 7 entirely, and that we hope is much clearer.\n\nWe start with Defn. 1:\n$$ Λ_π^{λ_1,λ_2} := \\|\\| Q_π^{λ_1} - Q_π^{λ_2}\\|\\| = \\Big\\|\\Big\\|W\\Big(I-γTK_π^{λ_1})^{-1}-(I-γTK_π^{λ_2})^{-1}\\Big):R^{SA}\\Big\\|\\Big\\|,$$\nwhere $K^λ_π := (λ Π^S + (1-λ) \\Phi W^Π)$, and the `:` operator denotes a tensor product over two indices.\n\nLet us parse this equation by considering one of the two Q-functions in isolation (say $λ_1$).\n\n**Policy Tensors.** First, note that $Π^S$ and $\\Phi W^Π$ are mappings from states to distributions over state-action pairs $s\\mapsto\\Delta(\\tilde s,a)$. We call the former the *MC policy tensor*, ${Π^S}\\_{s,\\tilde s,a}=𝟙[s=\\tilde s]\\sum\\_ω \\Phi\\_{s,ω}π\\_{ω,a}$, which maps each state the expected policy under its observation distribution. We call the latter the *TD policy tensor*, $(\\Phi W^Π)_{s,\\tilde s,a}=\\sum\\_ω \\Phi\\_{s,ω} π\\_{ω,a} W\\_{ω,\\tilde s}$, which spreads the policy probabilities for a given observation $ω$ across all states $\\tilde s$ that produce $ω$, weighted by $W=Pr(s|ω)$. $K\\_π^{λ_1}$ is a convex combination of these policy tensors, parameterized by ${λ_1}$.\n\n**Transition-policy tensors.** Left-multiplying a policy tensor by $T$ produces an $(S×A×S×A)$ transition-policy tensor, which describes the probability of each state-action transition $(s,a)\\mapsto\\Delta(s',a')$ under the policy. Intuitively, TD($λ_1$) computes Q-values for a POMDP as though state-action pairs are evolving according to $TK\\_π^{λ_1} = λ_1 TΠ^S + (1-λ_1)T\\Phi W^Π$, which is a mixture of the policy dynamics under two models: the MC transition-policy ($TΠ^S$) and the TD transition-policy ($T\\Phi W^Π$).\n\n**Final projections.** The expression $(I - γTK\\_π^{λ_1})^{-1}$ computes the state-action successor representation for this hybrid transition model, then multiplication by $R^{SA}$ computes state-space $Q$-values. Finally, these are projected through $W$ to compute observation-space $Q$-values.\n\nReturning to the question of when $Λ=0$ for all policies, there are 3 cases:\n\n1. Policy tensors are equal: $K_{π}^{λ_1}=K_{π}^{λ_2}$. This simplifies to $Π^S=\\Phi W^Π$, which is satisfied if and only if $\\Phi W = I$, i.e. each observation is generated by exactly one hidden state. This means the POMDP is actually a block MDP.\n\n2. Transition-policy tensors are equal: $TK_{π}^{λ_1}=TK_{π}^{λ_2}$. In this case, the POMDP is equivalent to an MDP. Both transition-policy tensors lead to the same state-action successor representation, so any reward function will lead to the same value under both models.\n\n3. Any differences between transition-policy tensors (and thus the state-action successor representation) are subsequently projected away by either $W$, $R^{SA}$, or both. In this case, LD=0, but the effective MDP may be lossy. Fortunately, such pathological environments appear to be quite rare, and they are easy to handle, as we discuss below.\n\nParity Check is an example of **case 3**. One can check (by hand or using our code) that the $λ_1$ and $λ_2$ transition-policies (and thus the successor representations) differ with and without aliasing. However, these differences are perfectly symmetric with respect to observations and rewards, so they are projected away by $W$ and $R^{SA}$. The Q-values are therefore zero at every observation, and so the LD=0 for all policies. Nevertheless, as we cover in our global response, we can still detect and mitigate partial observability for this problem.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Thank you all for your feedback and suggestions. Your reviews helped us improve the quality of our submission.

We've made the following requested changes (**see PDF**):
1. Experiment varying “Markovness” of environment (new Fig 1)
2. Visual comparison of learned memories for PacMan (new Fig 2, right)
3. New results for Parity Check including λ-discrepancy (LD) detection (new Fig 3, left) and memory optimization (new Fig 3, right)
4. Architecture diagram to clarify training procedure and RNN results (new Fig 2, left)
5. Simplified Sec. 3, including discussion of edge cases (see response to reviewer **8yd3**)

We also justify our choice of environments and discuss selecting λ values below.

---

**Varying "Markovness" (pbjp, 8yd3)**

We ran an additional experiment to measure LD while interpolating between an MDP and a POMDP for the T-Maze environment. We interpolate between an identity observation function (in which each state leads to a distinct observation) and an aliased one (where certain states cannot be distinguished). In our experiment, the LD has a minimum at zero for MDP and increases with partial observability for three different types of state aliasing. The results are in new Fig 1, and we will add them to the camera-ready if the paper is accepted.

Note: this experiment necessarily compares LD between different environments, which it was not designed for. This can lead to counterintuitive results depending on discount factor and choice of weighted norm, since the size of the observation space is changing, and states from later in the trajectory become aliased with earlier ones. To mitigate these issues, (for this experiment only) we use the max norm and set γ≈1.

---

**Memory visualization**

We also added visualizations of the agent's memory functions. In PacMan, the food dots provide a convenient way to track the agent's path. By training a small neural network probe (offline) to predict the remaining food dot locations from the agent’s memory, we can visualize how well the agent remembers where it has been. We compare the memory of a baseline RNN agent against that of an RNN with LD minimization, and find that minimizing LD facilitates much more accurate reconstruction of the agent’s path. These visualizations are shown in new Fig. 2 (right).

---

**Parity Check (rpiT)**

While our analysis showed that symmetric environments like the Parity Check example are likely rare, our initial submission did not offer a practical mitigation strategy for them. However, it turns out such symmetries are actually easy to handle: we can simply add a small amount of (stochastic) memory. If the environment is truly an MDP, adding memory will have no effect on LD, but if it is a POMDP, a randomly initialized memory function may be enough to break symmetry and produce a LD>0. Indeed, almost all of the 10K random memory functions we generated led to substantial LD (Fig 3, left). A subsequent memory optimization (using the procedure in Sec. 4) was able to extract 85% of the normalized return from the environment (Fig 3, right). This helps explain why such edge cases are not an issue for our RNN experiments: random memory functions (like the ones we initialize with in Sec. 5) break the symmetries that would appear in POMDPs with LD=0.

---

**Environment Complexity (pbjp, 8yd3)**

Some reviewers argued that the environments we tested were “simple” or “low complexity”. Our objective was to test complexity in the required memory function. For this reason, we intentionally avoided high-dimensional observation spaces that would require substantial representation learning on top of the memory learning task we were investigating. While the observation spaces of the domains in Sec. 5 may seem “simple”, the memory functions required to solve these domains are very complex, as reviewer 8yd3 pointed out, “I appreciate the experiments on larger POMDP planning benchmarks, since they require complex memory that seem non-trivial to learn.” For example, in P.O. PacMan, the agent must remember its current position, track which food dots and power-ups it has consumed, and remember in which directions it has recently seen ghosts. This is a highly challenging memory task and is sufficient to make our scientific point, since baseline RNNs struggle to learn a useful memory function (see new Fig. 2, right). We agree that investigating larger domains is important for future work, but we do not want to obscure the point we make here by conflating memory learning with arbitrary feature learning.

---

**Understanding and selecting λ parameters (pbjp, AeCz, 8yd3)**

Recall that TD(λ) is a weighted average of n-step returns, where λ controls the weighting scheme. Including a bootstrapped value term in the return target ($… + γ^n V(ω_{t+n})$) implicitly treats the environment like an MDP, as if the observation $ω_{t+n}$ is a Markovian state. Low values of λ give more weight to shorter rollouts where the bootstrapped term is discounted less (and thus has greater significance), whereas larger λ values put more weight on longer rollouts where the bootstrapped term is more heavily discounted. With two distinct λs, we can see if the estimators diverge when they use bootstrapping components of differing strength (which does not happen for MDPs).

Our theoretical and empirical results indicate that $λ_1$ and $λ_2$ should be well separated to provide sufficient signal for learning memory. Indeed, our hyperparameter sweep in Sec. 5 selected λ parameters that reflect this intuition, with a large difference between $λ_1$ and $λ_2$ values. In principle, LD is maximized for λ in {0,1}, but in practice, setting λ too close to 0 can lead to high bias, especially during the early part of learning, and setting λ too close to 1 can lead to high variance in returns and instability in the estimator. This is why our experiments only test λ values that are in the range [0.1, 0.95].

We will add this discussion to the paper if it is accepted.

Meta-review: 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.
Decision: Accept (poster)