Title: Mitigating Partial Observability in Sequential Decision Processes via the Lambda Discrepancy
Venue: NeurIPS 2024
Accepted: yes
Keywords: reinforcement learning

partial observability

value estimation

memory
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

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.

1 Introduction

The dominant modeling frameworks for reinforcement learning [Sutton and Barto, 2018] define environments in terms of an underlying Markovian state representation. This modeling choice, called the Markov assumption , is nearly ubiquitous in reinforcement learning, because it allows environment dynamics, rewards, value functions, and policies all to be expressed as functions that are independent of the past given the most recent state. In principle, an environment can be modeled as either a Markov decision process (MDP) [Puterman, 1994], or its partially observable counterpart, a POMDP

[Kaelbling et al., 1998], as long as the underlying state representation contains enough information to satisfy the Markov assumption. The POMDP framework is more general, but POMDPs are typically much harder to solve than MDPs [Zhang et al., 2012], so it is important to know when it is appropriate to use the simpler MDP framework.

Ideally, a system designer can ensure that their reinforcement learning agent is configured to use the appropriate problem model. If an environment is partially observable, the designer may manually add relevant decision-making features to the agent's state representation, either by concatenating observations together [Mnih et al., 2015] or by using other types of feature engineering [Bellemare et al., 2020, Galataud, 2021, Tao et al., 2023]. Alternatively, the designer can manually specify a set of possible environment states over which the agent can maintain a Markovian belief distribution [Kaelbling et al., 1998]. The challenge is that it is not always obvious when the designer has supplied sufficient information to satisfy the Markov assumption. These approaches require the person deploying the agent to know details about the very problem the agent is supposed to help them solve.

An alternative approach-that we explore in this paper-is to let the agent learn a good state representation for the problem it is solving. The conventional deep-learning-era wisdom is that the best problem representations come from training a system to solve a task, rather than from human designers. When faced with a potentially partially observable environment, if we provide the agent with a large enough recurrent neural network (RNN), it can perhaps discover an internal state representation 'endto-end' that maximizes reward via gradient descent [Bakker, 2001, Hausknecht and Stone, 2015, Ni et al., 2022, Dong et al., 2022]. Indeed, end-to-end RNN architectures work well for many problems and are both general-purpose and scalable. However, these techniques only implicitly address the problem of learning a Markovian state representation; in this paper, we show that we can achieve much better learning performance when we explicitly tackle the problem of partial observability.

In our approach, the agent learns a state representation by directly minimizing a metric that assesses whether the environment is fully or partially observable. We call our metric the λ -discrepancy , and define it as the difference between two distinct value functions estimated using temporal difference learning (TD) [Sutton, 1988]. Specifically, the TD( λ ) method defines a smooth trade-off between one-step TD (i.e., λ = 0 ), which makes an implicit Markov assumption, and Monte Carlo (MC) estimation ( λ = 1 ), which does not, and intermediate λ values interpolate between these extremes. By comparing value estimates for two distinct values of λ , we can check that the agent's observations support Markovian value prediction, and augment them with memory if we find they are incomplete.

Our main contributions 4 are as follows:

We introduce and formally define the λ -discrepancy and prove that it is exactly zero for MDPs and almost always non-zero for a broad class of POMDPs that we characterize. This analysis tells us that our metric reliably detects partial observability.

We then consider a tabular, proof-of-concept experiment that adjusts the parameters of a memory function to minimize the λ -discrepancy via gradient descent. For this experiment, we compute the λ -discrepancy in closed form. Our results demonstrate that minimizing the λ -discrepancy is a viable path to mitigating partial observability.

Finally, we integrate our approach into a deep reinforcement learning agent and evaluate on a set of large and challenging POMDP benchmarks. We find that minimizing the λ -discrepancy between two learned value functions, as an auxiliary loss alongside traditional reinforcement learning, is often significantly more effective (and never worse) than training the same agent with just a single value function.

Overall, we find that the λ -discrepancy is a simple yet powerful metric for detecting and mitigating partial observability. The metric is also practical, since it can be computed directly from value functions, which are ubiquitous in reinforcement learning. Furthermore, such value functions need only consider observable parts of the environment, so the λ -discrepancy remains applicable even without the common assumption that the agent knows the full set of possible POMDP states.

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 R : S×A → R , transition function T : S×A → ∆ S mapping to a distribution over states, discount factor γ ∈ [0 , 1] , and initial state distribution p 0 ∈ ∆ S . The agent's goal is to find a policy π S : S → ∆ A that selects actions to maximize return , g t , the discounted sum of future rewards starting from timestep t : g π S t = ∑ ∞ i =0 γ i r t + i , where r i is the reward at timestep i . We denote the expectation of these returns as value functions V π S ( s ) = E π S [ g t | s t = s ] and Q π S ( s, a ) = E π S [ g t | s t = s, a t = a ] .

We are interested in the learning setting, where the agent has no knowledge of the underlying state s nor even the set of possible states S (let alone the transition, reward, and observation functions). It receives an observation ω t ∈ Ω at each timestep and must find a way to maximize expected return. One way to do this is to construct a state representation , perhaps using some form of memory, on which it can condition its behavior. A state representation z ∈ Z is Markovian if at any timestep t , the representation z t and action a t together are a sufficient statistic for predicting the reward r t and next representation z t +1 , instead of requiring the agent's whole history:

The definition can be applied to states ( s ), observations ( ω ), and even memory-augmented observations (defined in Sec. 4), by replacing z with the relevant quantity. States and observations are equivalent in MDPs, and this property is satisfied for both by definition, but in POMDPs it typically only holds for the underlying, unobserved state s -not the observations.

Markovian state representations have several desirable implications. First, if the Markov property holds then so does the Bellman equation: V π S ( s ) = ∑ a ∈A π S ( a | s ) ( R ( s, a ) + γ ∑ s ′ ∈S T ( s ′ | s, a ) V π S ( s ′ ) ) . The Bellman equation allows agents to estimate the value of a policy, from experiences, and without knowing T or R , via a recurrence relation over one-step returns,

which converges to the unique fixed point V π S . A second implication is that the transition and reward functions, and consequently the value functions V π S and Q π S , have fixed-sized inputs and are therefore easy to parameterize, learn, and reuse. Finally, it follows from the Markov property that the optimal policy π ∗ S can be expressed deterministically and does not require memory [Puterman, 1994].

We can unroll the Bellman equation over multiple timesteps to obtain a similar estimator that uses n -step returns: V π S ( s ) = E π S [ g t : t + n | s t = s ] , where g t : t + n := r t + γr t +1 + γ 2 r t +2 + · · · + γ n V π S ( s t + n ) , with g t : t + n := g t if the episode terminates before t + n has been reached. The same process works for weighted combinations of such returns, including the exponential average:

with V λ =1 π S ( s ) = E π S [ g t | s t = s ] as a special case. Equation (3) defines the TD( λ ) value function as an expectation over the so-called λ -return [Sutton, 1988]. Low values of λ give more weight to shorter returns where the value term V π S ( s t + n ) is discounted less (and thus has greater significance), whereas larger λ values put more weight on longer returns where the value term is more heavily discounted. Given an MDP and a fixed policy, the recurrence relations for all TD( λ ) value functions share the same fixed point for any λ ∈ [0 , 1] ; however, if the Markov property does not hold, different λ may have different TD( λ ) fixed points. In this work, we seek to characterize this phenomenon and leverage it for detecting and mitigating partial observability.

+1

K

个

+1

3 Detecting Partial Observability

Before introducing our partial observability metric, let us first consider the T-maze example of Figure 1 under our two candidate decision-making frameworks: MDPs and POMDPs. In the T-maze, the agent only observes the color of its current square, and must remember the start square color (sampled uniformly from BLUE / RED ) to select the action at the junction that leads to the positive reward.

To model the T-maze as a POMDP, the state must include the agent's position and the goal location. The transitions are deterministic; each action moves the agent in the corresponding direction. The goal location is sampled at the start of an episode, and defines both the reward function and the observation for the initial square. The observation function uniquely identifies these starting states, but all corridor states are aliased together (they map to the same observation), as are the two junction states.

We can convert our POMDP model into an MDP model by using observations Ω as the MDP's 'states'. This requires new transition and reward functions, T Ω and R Ω , which we define as: T Ω ( ω ′ | ω, a ) := ∑ s,s ′ ∈S Φ( ω ′ | s ′ ) T S ( s ′ | s, a ) Pr( s | ω ) and R Ω ( ω, a ) := ∑ s ∈S R S ( s, a ) Pr( s | ω ) , where Pr( s | ω ) is policy-dependent and describes how each hidden state s i ∈ S contributes to the overall environment behavior when we see observation ω . While there are several sensible choices of weighting scheme for this type of state aggregation (uniform over states, relative frequency, etc.), only one of these choices coincides with the distribution Pr( s | ω ) . That particular weighting scheme is the one that averages the time-dependent Pr( s t | ω t ) over all timesteps, weighted by visitation probability under the agent's policy, and discounted by γ . 5 We explain how to construct this expression in Appendix C. We call the MDP defined in this way the effective MDP model for a given POMDP.

The effective MDP model marginalizes over histories (and POMDP hidden states). For example, the model predicts that going UP from the junction will reach the goal exactly half the time, because the transition dynamics marginalize over (equally likely) starting observations. Note that the POMDP model does no such averaging: if the agent initially observes BLUE , then going UP from the junction will always reach the goal, but DOWN never will. Thus, we see that, for the T-maze, there is a mismatch between the POMDP model and the effective MDP model. In other words, the POMDP's hidden states S are a Markovian representation, but its observations Ω are not, despite the fact that the effective MDP model treats them as Markovian 'states'.

In principle, an agent could measure partial observability by simultaneously modeling its environment as both a POMDP and an MDP, and comparing the models' predictions. Unfortunately, since the agent lacks any information about the unobserved state space S , the POMDP model would require variable-length history inputs at each time step, and would come with computational and memory costs that grow exponentially with history length. Instead, we propose a model-free approach, using value functions conditioned only on observations, to approximate this comparison in a tractable way.

3.1 Value Function Estimation under Partial Observability

The Bellman equation and its sample-based recurrence relations (2) and (3) are defined for Markovian states. If we apply them to the observations of a POMDP, we are actually working with the effective MDP model of that POMDP, instead of the POMDP itself. To see this, consider one-step TD ( λ = 0 ),

Memory

π

PPO

Action

RNN

VTD

TD

A =IIvTD _ vMCI

Observation

VMC

MC

LActor

r=LPPO

LCritic = β·A+ (1 -β)·(CvrD + LvMc)

where we use the same recurrence relation (2) but now our expectation is sampling from the POMDP:

where we have suppressed V Ω 's dependence on the observation-based policy π for ease of notation.

We see from Equation (4) that the value function TD computes for a POMDP is the fixed point of the Bellman operator for the effective MDP model. 6 By contrast, the Monte Carlo value estimator ( λ = 1 ) does not exploit the Bellman equation at all; it simply averages over returns: V λ =1 π ( s ) = E π [ g t | s t = s ] . Translating the Monte Carlo estimator to POMDPs merely requires one additional expectation to convert from states to observations:

This means MC ( λ = 1 ) effectively takes the hidden-state value function V S of the POMDP and projects it into observation space, whereas TD ( λ = 0 ) directly computes the value function for the projected model as an MDP, by treating observations as states. Interpolating λ between 0 and 1 smoothly varies the value estimate's reliance on the effective MDP model.

For generality, we derive an expression (see Appendix C) for Q λ π -values in terms of a given λ parameter (expressed in tensor notation for compactness) to reveal how λ trades off between the projected state value function and the value function of the projected model:

where the tensor product . . contracts two indices instead of one, 7 with tensors defined as follows:

λ

Q π ( Ω ×A ) is a matrix of Q -values; W ( Ω ×S ) contains the state-blending weights given by Pr( s | ω ) for observation ω ; I ( S × A × S × A ) is an identity tensor with I sas ′ a ′ = 1 [ s = s ′ ] 1 [ a = a ′ ] ; T ( S × A × S ) contains the hidden-state transition probabilities T ( s ′ | a, s ) ; Π S ( S × S × A ) contains the policy spread over hidden states (see below); Φ ( S × Ω ) is the observation function, containing probabilities Φ( ω | s ) ; W Π ( Ω ×S × A ) combines W with Π S to obtain probabilities Pr( s, a | ω ) ; R S A ( S × A ) contains the hidden-state rewards R ( s, a ) .

Let us take a moment to parse this equation. First, Π S and Φ W Π are mappings from states to stateaction pairs. We call the former the MC policy spread , Π S s,s ′ ,a = 1 [ s = s ′ ] ∑ ω Φ s,ω π ω,a , which maps states to the expected policy under their observation distribution. We call the latter the TD policy spread , (Φ W Π ) s,s ′ ,a = ∑ ω Φ s,ω π ω,a W ω,s ′ , which reallocates the policy probabilities for a given observation ω across all states s ′ that produce that observation, weighted by W . K λ π is a convex combination of these two policy spread tensors, parameterized by λ . Multiplying a policy spread tensor on the left by T produces an ( S×A×S×A ) transition-policy tensor, describing the probability of each state-action transition ( s, a ) ↦→ ∆( s ′ , a ′ ) under the policy. Intuitively, Equation (6) says that TD( λ ) computes Q-values for a POMDP as though state-action pairs are evolving according to TK λ π = λT Π S +(1 -λ ) T Φ W Π , which is a mixture of the policy dynamics under two transition models: the MC transition-policy ( T Π S ) and the TD transition-policy ( T Φ W Π ). The expression ( I -γTK λ π ) -1 . . R computes the state-space Q -values for this hybrid transition model. Finally, these are projected through W to compute the observation-space Q -values. (See Appendix C for more details.)

3.2 λ -Discrepancy

We have shown that, under partial observability, Q λ π value functions may differ for different λ parameters, due to varying reliance on the effective MDP model in the TD( λ ) estimator. We call this difference the λ -discrepancy , and we propose to use it as a measure of partial observability.

No Aliasing

Junction Aliased

0

0.6

0

0

0

3

3

3

3

State

6

0.4

Aliasing Type

State

6

State

9

State

9

9

Corridor

9

9

9

12

0.2

Junction

12

12

12

Both

0

36912

0.0

0

36912

0

36912

036912

Observation

0.0

0.2

0.4

0.6

0.8

1.0

Observation

Observation

Observation

Mixing Parameter (p)

ΦPerfect

ΦAliased

Definition 1 For a given POMDP model P and policy π , the λ -discrepancy Λ λ 1 ,λ 2 P ,π is the weighted norm of the difference between the Q -functions estimated by TD( λ ) for λ ∈ { λ 1 , λ 2 } :

The choice of norm can be arbitrary, as can the norm weighting scheme, as long as it assigns positive weight to all reachable observation-action pairs. We discuss choices of weighted norm in Appendix G.2. For brevity, we suppress the P subscript when the POMDP model is clear from context.

Auseful property (that we will prove in Theorem 2) is that if the POMDP has Markovian observations, the λ -discrepancy is exactly zero for all policies. However, for it to be a useful measure of partial observability, we must also show that the λ -discrepancy is reliably non-zero when observations are non-Markovian. For this, we have the following theorem.

̸

Theorem 1 Given a POMDP model P and distinct λ = λ ′ , if there exists a policy π : Ω → ∆ A such that Λ λ,λ ′ P ,π = 0 , then Λ λ,λ ′ P ,π = 0 for all policies except at most a set of measure zero.

̸

̸

Proof sketch: We formulate the λ -discrepancy of a given policy as the norm of an analytic function, then use the fact that analytic functions are zero everywhere or almost nowhere, along with the fact that norms preserve this property. The full proof is given in Appendix D.

Intuitively, Theorem 1 says that the λ -discrepancy can detect non-Markovian observations. If it is possible to reveal that a POMDP's observation-space value function does not match that of its effective MDP model, then almost all policies will do so. The theorem further suggests that even if a particular policy has zero λ -discrepancy, small perturbations to that policy will almost surely detect partial observability if it is present.

We also find that increasing amounts of partial observability lead to larger λ -discrepancy. In Figure 3 we interpolate between Markovian and aliased T-Maze environments, and consider three types of state aliasing. Here we use a fixed policy that goes RIGHT until the junction, then UP with probability 2 / 3 and DOWN w.p. 1 / 3 . To avoid artifacts due to interactions between the discount factor, weighted norm, and changing observation function, we set γ ≈ 1 and use the (policy-weighted) maximum norm (see Appendix G.2). 8 The results confirm that the λ -discrepancy is a useful indicator of partial observability, provided that it is non-zero for at least one policy.

3.3 What conditions cause the λ -discrepancy to be zero?

We can characterize the cases in which the λ -discrepancy is zero for all policies by inspecting Definition 1. Because norms are positive definite, it suffices to consider the expression inside the norm. The only ways for this expression to equal zero are either when the difference between policy spread tensors ( K λ 1 π -K λ 2 π ) is zero (which we will show implies Markovian observations), or it

→

+1

个

O

K

↑个

R

O

个

O

-1

→

个个

O

K

→

O

+1

ofmemoryfunctions

Parity Check with Random Memory

NormalizedReturn

1.0 ---=

400

0.5

0.851

200

0.000

No.

0

0.0

10-7

10-6

10-5

10-4

10-3

10-2

0

1

X-discrepancy (A)

MemorySize(bits)

is non-zero but is projected away by T , R SA , and/or W . We first consider when the policy spread tensors-which are the only terms that depend on λ -are equal, i.e. K λ 1 π = K λ 2 π .

̸

Theorem 2 For any POMDP P and any λ 1 = λ 2 , K λ 1 π = K λ 2 π if and only if P is a block MDP.

Proof sketch: Using Eq. (6), K λ 1 π = K λ 2 π can be simplified to Π S = Φ W Π , which is satisfied if and only if Φ W = I , i.e. each observation is generated by exactly one hidden state. Proof in Appendix E.

Thus, when the policy spread tensors are equal, it means there is no state aliasing and P is a block MDP. Now consider the cases where the difference between K λ π is projected away by T , R SA , and/or W , starting from the policy spread tensors in Definition 1 and working our way outwards.

If the policy spread tensors differ, the transition tensor may project these differences away ( TK λ 1 π = TK λ 2 π ), for example, if the transition probabilities leading to the aliased states always occur in some fixed proportion. In such cases, we can collapse the aliased states to non-aliased ones without any loss of information by reformulating the incoming transition probabilities to the aliased states as stochastic transition dynamics (and possibly rewards) after the non-aliased states. In other words, P is equivalent to an MDP. We prove this claim and provide an example of such an environment in Appendix F.

Finally, differences between transition-policies may be projected away by the state-blending weights W , the reward function R SA , or both. W ( I -γTK λ 1 π ) -1 . . R SA = W ( I -γTK λ 2 π ) -1 . . R SA . In this case, the effective MDP may be lossy, which is a limitation of the λ -discrepancy metric. Fortunately, such environments appear to be rare, and they are easy to handle, as we show in the example below.

Parity Check Environment. The Parity Check POMDP of Figure 4 has four equally likely starting states, each associated with a pair of colors that the agent will see during the first two timesteps. At the subsequent (white) junction state, the agent must decide based on these colors whether to go UP (black arrow) or DOWN (white arrow). The rewards are defined such that UP is optimal if and only if the color family matches (i.e. RED → PINK or BLUE → CYAN ).

The transition-policies for the first eight states differ with and without aliasing, but the differences are perfectly symmetric with respect to rewards and observations. 9 Every observation has zero expected return, and so the λ -discrepancy is zero for all policies. Our analysis in Appendix H.5 suggests such edge-cases are the exception and not the rule. Even so, such examples are easy to handle: we can simply add a small amount of memory. If the environment is truly an MDP, adding memory should have no effect, but if it is a POMDP, a randomly initialized memory function (defined in the next section) may be enough to break symmetry and produce a λ -discrepancy (as seen in Figure 4, center).

Taken together, Theorems 1 and 2 suggest that the λ -discrepancy could be a useful measure for detecting partial observability. In the next section, we demonstrate the efficacy of using the λ -discrepancy to learn memory functions that mitigate partial observability.

4 Memory Learning with the λ -Discrepancy

The λ -discrepancy can identify that we need memory, but it cannot yet tell us what to remember. For that, we must replace the POMDP P in Definition 1 with a memory-augmented version. In general,

PerformancewithMemoryOptimization

Return

1.0

MemorySize

(bits)

0.8

orma

0.6

3

0.4

T-maze

Paint

Network

Tiger

4x3

Shuttle

Cheese

a memory can be any mapping from agent histories ( ω 0 , a 0 , ..., a t -1 , r t -1 , ω t ) to internal memory states m within some set of memory states M . For practical reasons, we restrict our attention to recurrent memories that update an internal representation incrementally from fixed-size inputs.

We define a memory function µ : Ω × A × M → M as a mapping from an observation, action, and memory state, ( ω, a, m ) , to a next memory state m ′ = µ ( ω, a, m ) . Given a POMDP P , a memory function µ induces a memory-augmented POMDP P µ , with extended states S M = S × M , actions A M = A×M , and observations Ω M = Ω ×M . The augmented transition dynamics T M preserve the original transition dynamics T for states S while allowing the agent full observability and control over memory states M . Memory functions naturally lead to memory-augmented policies π µ : Ω M → ∆ A M and value functions V π,µ : Ω M → R and Q π,µ : Ω M ×A M → R that reflect the expected return under such policies. For details, see Appendix G.1.

The λ -discrepancy (Definition 1) applies equally well to memory-augmented POMDPs, and can thus be used to determine whether a particular memory function µ induces a POMDP P µ with a Markovian observation space Ω M . We can also use it as a training objective for learning such a memory function. We conduct a proof-of-concept experiment on several classic POMDPs for which we can obtain closed-form gradients of the λ -discrepancy with respect to a parametrized memory function. In each domain, we randomly generate a set of stochastic policies, select the one with maximal λ -discrepancy Λ 0 , 1 P , and adjust the parameters of a memory function µ to minimize Λ 0 , 1 P µ via gradient descent. Figure 5 shows the improvement in policy gradient performance due to the resulting memory function for increasing memory sizes. The details of this experiment are provided in Appendix H. We also run the same experiment on the Parity Check example, and provide results in Figure 4 (right).

Wesee that the λ -discrepancy can help mitigate partial observability, but it is somewhat inconveniently defined in terms of the closed-form value function fixed-points. Fortunately, with the appropriate choice of weighted norm, we can estimate the λ -discrepancy purely from sampled observation-action pairs (or, in the case of memory, observation-memory-action tuples):

The norm is taken over the on-policy joint observation-action distribution Pr( ω ) π ( a | ω ) , which allows us to estimate the metric using samples generated by the agent's interaction with the environment. We show how to use it as an optimization objective in the following section.

5 A Scalable, Online Learning Objective

So far, we have shown that the λ -discrepancy can detect partial observability in theory and can mitigate it under certain idealized conditions. Now we demonstrate how to integrate our metric into samplebased deep reinforcement learning to solve problems requiring large, complex memory functions.

5.1 Combining the λ -Discrepancy with PPO

To minimize the λ -discrepancy, we augment a recurrent version of the proximal policy optimization (PPO) algorithm [Schulman et al., 2017] with an auxiliary loss. We use recurrent PPO as our base algorithm due to its strong performance in many POMDPs [Ni et al., 2022], and since the λ -discrepancy is a natural extension of generalized advantage estimation [Schulman et al., 2016], which

RockSample 11 × 11

RockSample 15 × 15

RNN +LD

Episode Reward

20

20

0

0

0

1

2

3

0.0

0.2

0.4

0.6

0.8

1.0

Timestep

×106

Timestep

×107

Battleship 10 × 10

Partially Observable PacMan

Episode Reward

15

2000-

10

1000

PPO+RNN+LD

PPO+RNN

PPO

5

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Timestep

×107

Timestep

×107

RNN

is used in PPO. In this algorithm, a recurrent neural network [Amari, 1972] (specifically a gated recurrent unit, or GRU [Cho et al., 2014]) is used as the memory function µ that returns a latent state representation given previous latent state and an observation. This latent representation is used as input to an actor network to return a distribution over actions, as well as a critic. The critic is usually a value function network which learns a truncated TD( λ ) value estimate for its advantage estimate.

To estimate the λ -discrepancy, we learn two TD( λ ) value function networks with different λ , parameterized by θ V, 1 and θ V, 2 respectively, and minimize their mean squared difference as an auxiliary loss:

where z t = µ θ RNN ( ω t , z t -1 ) is the latent state output by the GRU, and θ represents all parameters. We train this neural network end-to-end with the standard PPO actor-critic losses and backpropagation through time [Mozer, 1995]. Full details of the algorithm are provided in Appendix I.1.

Any algorithm that uses value functions could, in principle, also leverage the λ -discrepancy. However, the theoretical results in Section 3 only hold for value function fixed points and not necessarily for estimated values prior to convergence. Inaccurate value functions may lead to an inaccurate λ -discrepancy. Fortunately, value functions often provide a useful learning signal before convergence, and we observe in the following sections that the λ -discrepancy is also useful during learning.

5.2 Large Partially Observable Environments

We evaluate our approach on a suite of four hard partially observable environments that require complex memory functions. Battleship [Silver and Veness, 2010] requires reasoning about unknown ship positions and remembering previous shots. Partially observable PacMan [Silver and Veness, 2010] requires localization within the map while tracking dot status and avoiding ghosts with only short-range sensors. RockSample (11, 11) and RockSample (15, 15) [Smith and Simmons, 2004] have stochastic sensors and require remembering which rocks have been sampled. While these environments were originally used to evaluate partially observable planning algorithms in large-scale POMDPs [Silver and Veness, 2010], we use them to test our sample-based learning algorithms due to the complexity of their required memory functions. 10 See Appendix I.2 for more details.

5.3 Experiments

We conduct experiments with regular PPO, recurrent PPO, and our λ -discrepancy-augmented recurrent PPO in Figure 6. We also visualize the agent's memory for the partially observable PacMan environment by reconstructing the dot locations from the RNN latent state to show where the agent 'thinks' it has been (see Appendix I.5). We performed a hyperparameter sweep for each method and report learning curves for undiscounted return (see Appendix I.3 for discounted learning curves and Appendix I.4 for additional experimental details).

The λ -discrepancy objective leads to significantly better final performance and learning rate versus recurrent and memoryless PPO in all environments. In RockSample (15 , 15) , the baseline agents quickly learn to exit immediately (for +10 reward), but never improve on this. By contrast, minimizing λ -discrepancy helps the agent learn the missing features that allow it to express better policies. We also run experiments on the classic POMDPs from Section 4, but due to the size of these problems, both baseline and our proposed approach performed almost optimally (see Appendix I.6).

We also observe that the best performing hyperparameters from our sweep offer further evidence for the theory developed in Section 3. Our theory suggests the λ parameters should be well separated in order to see the largest λ -discrepancy. Indeed, we find that a large difference between λ 1 and λ 2 was beneficial, with optimal λ s close to either 0 or 1 . In addition to this, our hyperparameter sweep also includes a PPO variant with two different TD( λ ) value functions but without the λ -discrepancy auxiliary loss described in Section 5.1. These agents were never selected in the sweep; using the loss in Equation 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 et al. [2021] investigated the state abstraction perspective and characterized the properties under which abstract state representations of MDPs either do or do not preserve the Markov property. Several other approaches characterize and measure partial observability and POMDPs. While POMDPs have been shown to be computationally intractable in general [Papadimitriou and Tsitsiklis, 1987], various works have studied complexity measures [Zhang et al., 2012] and defined subclasses with tractable solutions [Littman, 1993, Liu et al., 2022].

The most common strategies for mitigating partial observability are memory-based approaches that summarize history. Early approaches relied on discrete representations of history, including tree representations [McCallum, 1996] or finite-state controllers [Meuleau et al., 1999]. Modern approaches mostly use RNNs [Amari, 1972] trained via backpropagation through time (BPTT) [Mozer, 1995] to tackle non-Markovian decision processes [Schmidhuber, 1990]. Various approaches use recurrent function approximation to learn better state representations. One successful approach is learning a recurrent value function [Lin and Mitchell, 1993, Bakker, 2001, Hausknecht and Stone, 2015] that uses TD error as a learning signal for memory. Policy gradient methods, including PPO (which we compare to), have also been used with recurrent learning to mitigate partial observability [Wierstra et al., 2007, Heess et al., 2015]. Model-based methods can learn a recurrent dynamics model [Hafner et al., 2020] to facilitate planning alongside reinforcement learning. These approaches learn their representations implicitly to improve prediction error, rather than explicitly to mitigate partial observability.

7 Conclusion

We introduce the λ -discrepancy: an observable and differentiable measure of non-Markovianity suitable for mitigating partial observability. The λ -discrepancy is the difference between two distinct value functions estimated using TD( λ ), for two different λ values. We characterize the λ -discrepancy and prove that it reliably distinguishes MDPs from POMDPs. We then use it as a memory learning objective and demonstrate that minimizing it in closed-form helps learn useful memory functions in small-scale POMDPs. Finally, we propose a deep reinforcement learning algorithm which leverages the λ -discrepancy as an auxiliary loss, and show that it significantly improves the performance of a baseline recurrent PPO agent on a set of large and challenging partially observable tasks.

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-concept experiments. AK led the theoretical work, with support from CA, SL, RP, and RYT. RYT led 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 λ -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 to Aaron Kirtland and Sam Lobel, ONR grant N00014-22-12592, a gift from Open Philanthropy to the Center for Human-Compatible AI at Berkeley, and an AI2050 Senior Fellowship for Stuart Russell from the Schmidt Fund for Strategic Innovation.

References

Cameron Allen, Neev Parikh, Omer Gottesman, and George Konidaris. Learning Markov state abstractions for deep reinforcement learning. In Advances in 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. IEEE Transactions on Computers , C-21(11):1197-1206, 1972.

Bram Bakker. Reinforcement learning with long short-term memory. In Advances in Neural Information Processing Systems , volume 14, 2001.

Marc G. Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C. Machado, Subhodeep Moitra, Sameera S. Ponda, and Ziyu Wang. Autonomous Navigation of Stratospheric Balloons using Reinforcement Learning. Nature , 588(7836):77-82, 2020.

Cl´ ement Bonnet, Daniel Luo, Donal John Byrne, Shikha Surana, Sasha Abramowitz, Paul Duckworth, Vincent Coyette, Laurence Illing Midgley, Elshadai Tegegn, Tristan Kalloniatis, Omayma Mahjoub, Matthew Macfarlane, Andries Petrus Smit, Nathan Grinsztajn, Raphael Boige, Cemlyn Neil Waters, Mohamed Ali Ali Mimouni, Ulrich Armel Mbou Sob, Ruan John de Kock, Siddarth Singh, Daniel Furelos-Blanco, Victor Le, Arnu Pretorius, and Alexandre Laterre. Jumanji: a diverse suite of scalable reinforcement learning environments in JAX. In 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 Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

Anthony R. Cassandra. The POMDP page. https://www.pomdp.org/ , 2003.

Anthony R. Cassandra, Leslie P. Kaelbling, and Michael L. Littman. Acting optimally in partially observable stochastic domains. In Proceedings of the Twelfth AAAI Conference on Artificial Intelligence , pages 10231028, 1994.

Kyunghyun Cho, Bart van Merri¨ enboer, 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 in Natural Language Processing (EMNLP) , pages 1724-1734, 2014.

Lonnie Chrisman. Reinforcement learning with perceptual aliasing: The perceptual distinctions approach. In Proceedings of the Tenth AAAI Conference on Artificial Intelligence , pages 183-188, 1992.

Shi Dong, Benjamin Van Roy, and Zhengyuan Zhou. Simple agent, complex environment: Efficient reinforcement learning with agent states. Journal of Machine Learning Research , 23:1-54, 2022.

Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning , pages 1665-1674. Proceedings of Machine Learning Research, 2019.

Antoine Galataud. HVAC processes control: Can AI help?, 2021.

Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on 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 NeurIPS , 2015.

Herbert Jaeger. Tutorial on training recurrent neural networks, covering bppt, rtrl, ekf and the echo state network approach. GMD-Forschungszentrum Informationstechnik, 2002. , 5, 2002.

Leslie P. Kaelbling, Michael L. Littman, and Anthony R. Cassandra. Planning and acting in partially observable stochastic domains. Artificial Intelligence , 101(1-2):99-134, 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. A primer of real analytic functions . Springer Science & Business Media, 2002.

Nicholas Kushmerick, Steve Hanks, and Daniel S. Weld. An algorithm for probabilistic planning. Artificial Intelligence , 76(1-2):239-286, 1995.

Lihong Li, Thomas J. Walsh, and Michael L. Littman. Towards a unified theory of state abstraction for MDPs. In Proceedings of the Ninth International Symposium on Artificial Intelligence and Mathematics , 2006.

Long-Ji Lin and Tom M. Mitchell. Reinforcement learning with hidden states. From animals to animats , 2: 271-280, 1993.

Michael L. Littman. An optimization-based categorization of reinforcement learning environments. In From Animals to Animats 2: Proceedings of the Second International Conference on Simulation of Adaptive Behavior . The MIT Press, 1993. ISBN 9780262287159.

Qinghua Liu, Alan Chung, Csaba Szepesvari, and Chi Jin. When is partially observable reinforcement learning not scary? In Proceedings of Thirty Fifth Conference on Learning Theory , volume 178 of Proceedings of Machine Learning Research , pages 5175-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 , 35:16455-16468, 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 a real analytic function. Matematicheskie Zametki , 107(3):473475, 2020.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature , 518:529-533, 2015.

Volodymyr Mnih, 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-1937, 2016.

Michael Mozer. A focused backpropagation algorithm for temporal pattern recognition. Complex Systems , 3, 1995.

Vinod Nair and Geoffrey E. Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th International Conference on International Conference on Machine Learning , page 807-814, 2010.

Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free RL can be a strong baseline for many POMDPs. In Proceedings of the 39th International Conference on 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):441-450, 1987.

Ronald Parr and Stuart Russell. Approximating optimal policies for partially observable stochastic domains. In Proceedings of the Fourteenth International Joint Conference on Artificial Intelligence - Volume 2 , page 1088-1094, 1995.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In Proceedings of the 30th International Conference on Machine Learning , volume 28, pages 1310-1318, 2013.

Martin L. Puterman. Markov Decision Processes . Wiley, 1994.

J¨ urgen Schmidhuber. Reinforcement learning in Markovian and non-Markovian environments. In Advances in Neural Information Processing Systems , volume 3, 1990.

John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel. High-dimensional continuous 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 in Neural Information Processing Systems , 2010.

Trey Smith and Reid Simmons. Heuristic search value iteration for POMDPs. In Proceedings of the 20th Conference on Uncertainty in Artificial Intelligence , page 520-527, 2004.

Richard S. Sutton. Learning to predict by the methods of temporal differences. Machine learning , 3(1):9-44, 1988.

Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction . A Bradford Book, 2018.

Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems , volume 12, pages 1057-1063, 1999.

Ruo Yu 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¨ urgen 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 POMDPplanning and learning. In Proceedings of the Twenty-Sixth AAAI Conference on Artificial Intelligence , 2012.

A Limitations

The main theoretical limitations of our work come from the class of POMDPs that the λ -discrepancy is relevant to. The λ -discrepancy is capable of detecting non-Markovian observations in a 'large' class of POMDPs, which we have described using an analytic almost-all characterization. However, there do exist POMDPs in which the fact that observations are non-Markovian cannot be detected by the λ -discrepancy, and we provide one such example in Section 3.2. The examples of such POMDPs we have found seem to exist on a 'knife's edge', meaning that slightly perturbing the parameters of the POMDP results in a non-zero λ -discrepancy. Furthermore, incorporating memory into the value function seems to reliably produce a λ -discrepancy, without having to modify the underlying decision process.

The recurrent deep reinforcement learning algorithm we develop is also limited in that it only approximates the λ -discrepancy. The auxiliary loss in Section 5 is only an approximation of the λ -discrepancy. It minimizes the difference in value function estimates as opposed to the difference in fixed points that we consider in Section 3. The loss function is therefore a departure from our theoretical analysis, though the empirical results suggest that the value functions need not have fully converged for the discrepancy to be useful. Future work ought to investigate why this works so well in practice.

B Broader Impact

As a work studying the foundations of partially observable reinforcement learning, this paper does not have any particular negative societal applications. Application domains could include robotics and tasks such as navigation. The risk for misuse of this work are no different than any other works aimed towards algorithmic improvements for reinforcement leanring. Negative impact mitigation strategies are likewise the same as other similar works.

C TD( λ ) Fixed Point

Here we derive the fixed point of the TD( λ ) action-value update rule in a POMDP, following the Markov version by Sutton [1988]. First, define the expected return given initial observation ω 0 and initial action a 0 as

We can define the n -step bootstrapped update rule from this given a value matrix Q π by replacing part of the term with coefficient γ n with a Q π value, e.g. for the n = 2 case, we get

Translating these expressions into matrix notation, we have

where the terms W , T , R , and Φ are as in Equation 6, and π is the Ω → ∆ A policy written as an Ω × A tensor with entries in [0 , 1] . In particular, W ω,s = Pr( s | ω ) , which averages Pr( s t | ω t ) over all timesteps, weighted by visitation probability and discounted by γ . This is a well-defined stationary quantity, and it can be computed as follows. First solve the system Cx = b to find the discounted state occupancy counts x = c ( s ) , where b = p 0 is the initial state distribution over S , and C = ( I -γ ( T π ) ⊤ ) accounts for the policy-dependent state-to-state transition dynamics T π defined by T π s,s ′ = ∑ ω ∑ a Pr( ω | s ) Pr( a | ω ) Pr( s ′ | s, a ) . Then Pr( s | ω ) ∝ c ( s ) ∗ Φ( ω | s ) , so we can just multiply these terms together and renormalize. For the 2-step bootstrapped update rule above, we have:

Let Π be an Ω × Ω ×A representation of the Ω ×A policy π with Π ω,ω ′ ,a = 1 [ ω = ω ′ ] π ω,a . Likewise, let Π S be the effective policy over latent states, an S×S×A representation of the matrix Φ π , i.e. Π S s,s ′ ,a = 1 [ s = s ′ ](Φ π ) s,a . Given a matrix of action-values Q ( i ) π , the n -step update rule is:

where R SA is the matrix of reward values as described in Equation 6, and we use . . to denote double tensor contraction, while all other tensor products contract a single index.

We also have the standard definition of the TD( λ ) update rule as:

We are concerned with the fixed point of this update rule, which we refer to as Q λ π :

Separating this into a reward part with factor R SA and a value part with factor Q λ π , we find that the value part is

and for the reward part, we have the coefficients of R SA in the table below for values of n and k

where each term is multiplied by (1 -λ ) W in front. We can then see by summing over rows before columns that the reward part is:

So we rewrite Q λ π as follows:

Now let W A = W ⊗ I A , which is Ω ×A×S×A , and W Π = Π . . W A , which is Ω ×S × A . Here, ⊗ means the Kronecker product. This essentially repeats the W matrix |A| times to incorporate actions into the tensor. Note that for any S×A tensor G , ( W A . . G ) = WG . This is because ( W A . . G ) ωa = ∑ s,a ′ W A ωasa ′ G sa ′ , and the only nonzero terms in this sum are those such that a = a ′ . For these indices, W A ωasa ′ = W ωs , so ∑ s,a ′ W A ωasa ′ G sa ′ = ∑ s W ωs G sa = ( WG ) ωa .

Also, let F = ( I SA -λγT Π S ) -1 , which is S × A × S × A . Then we find:

At this point, we can subtract the second term on the right hand side from both sides, factor out Q λ π on the right, and multiply by ( I Ω A -(1 -λ ) γW A . . F . . T ΦΠ ) -1 on the left of both sides to obtain:

where the last equality follows from expanding the geometric series. Now we use the identity ( A -B ) -1 = ∑ ∞ k =0 ( A -1 B ) k A -1 to find:

To recap our previous definitions, W is an Ω ×S tensor, I is an S×A×S×A tensor, T is an S×A×S tensor, Π S is an S×S×A tensor, Φ is an S× Ω tensor, W Π is an Ω ×S×A tensor, and R SA is an S×A tensor.

Lastly, we briefly note that one can get the V λ π values by replacing W in the above equation with W Π and using a double contraction:

We can confirm that V λ π = ∑ a π ω,a Q ω,a , by rewriting the expression on the right as follows:

where B = ( I -γT ( λ Π S +(1 -λ )Φ W Π )) -1 . . R SA is an S × A tensor.

D Proof of Theorem 1 (Almost All)

In this section we prove Theorem 1, that there is either a λ -discrepancy for almost all policies or for no policies. Fix λ and λ ′ . Recall that we define the λ -discrepancy as follows:

where K λ π = λ Π S +(1 -λ )Φ W Π . Let Y be the largest open set in the space of stochastic Ω ×A matrices, considered as a subset of R | Ω | ( |A|1) . Now consider the λ -discrepancy as a function of the policy π . In other words, we define

Let X be an open subset of R n . We say that a function f : X → R is real analytic on X if for all x ∈ X , f can be written as a convergent power series in some neighborhood of x . For this proof, we will utilize the following facts:

the composition of analytic functions is analytic [Krantz and Parks, 2002],

the quotient of two analytic functions is analytic where the denominator is nonzero,

a real analytic function on a domain X is either identically zero or only zero on a set of measure zero [Mityagin, 2020].

We will also use the fact that for an invertible matrix A , each entry A -1 ij is analytic in the entries of A where the entries of A yield a nonzero determinant. We can prove this by first writing A -1 = det( A ) -1 adj( A ) = det( A ) -1 cof( A ) ⊤ where adj A is the adjugate of A and cof( A ) is the cofactor matrix of A . Each entry of the cofactor matrix is a cofactor that is polynomial in the entries of A , and is therefore analytic in them. Therefore, each entry of A -1 is the quotient of two analytic functions and is therefore analytic except where det A = 0 .

Next, we will show that Λ is an analytic function. Note that the variable terms in the equation are W , W Π , R SA , T , Π S , and Φ . Of these, R SA , T , and Φ are constant with respect to π . Π S ilj = ∑ k 1 [ i = l ]Φ ik π kj , so each entry of Π S is analytic on Y in the entries of π . Likewise, P ij = ∑ k,a Φ ik π ka T iaj is analytic on Y . Therefore, the state-occupancy counts c = p 0 + γP ⊤ p 0 + γ 2 ( P ⊤ ) 2 p 0 + · · · = ( I -γP ⊤ ) -1 p 0 , where p 0 contains the initial state probabilities, are the composition of analytic functions and thus analytic on Y . W ωs = Φ sω c s ∑ s ′ Φ s ′ ω c s ′ is analytic on Y for the same reason, and the denominator of W ωs , ∑ s ′ Φ s ′ ω c s ′ , is nonzero for all observations able to be observed with nonzero probability. Finally, Λ is then a composition of analytic functions on Y and thus analytic itself.

To handle the norm weighting, we note that w π is analytic in π as w π = (1 , π ( a | ω )) , and the dot product of w π with Λ is also analytic. Now, we use the fact mentioned above that the zero set of a nontrivial analytic function is of measure zero. Therefore, the zero set of Λ · w π is either zero for all policies or zero only on a set of measure zero. To finish, we note that because norms are positive definite, Λ π = 0 if and only if Λ · w π = 0 , so this result extends to the normed λ -discrepancy as well.

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 6 we define K λ π = λ Π S +(1 -λ )Φ W Π . Suppose K λ π = K λ ′ π . We can rewrite this as ( λ -λ ′ )Π S -( λ -λ ′ )Φ W Π = ( λ -λ ′ )(Π S -Φ W Π ) = 0 . This implies that either λ = λ ′ or Π S = Φ W Π .

Recall the definition of Π S and W Π in Section C. Expanding the equation Π S = Φ W Π using these definitions, we have that

Then by setting the rightmost expression in each equation equal and simplifying with the indicator function, we have that for all i, j, k ,

̸

̸

̸

We will first consider the case where i = j . We have that for all i = j and all k , ∑ ω Pr( ω | s i ) Pr( a k | ω )Pr( s j | ω ) = 0 . Because each term in the sum is nonnegative, this is equivalent to the statement that for all i = j , all k , and all ω , Pr( ω | s i )Pr( a k | ω )Pr( s j | ω ) = 0 . Now note that for all observations ω , there exists some k ′ such that Pr( a k ′ | ω ) > 0 . Therefore, we have that for all i = j and all ω , there exists a k ′ such that Pr( a k ′ | ω ) > 0 . This implies that for all i = j and all ω , Pr( ω | s i ) Pr( a k ′ | ω )Pr( s j | ω ) = 0 and thus Pr( ω | s i )Pr( s j | ω ) = 0 . This means that if state s i produces an observation ω , then ω cannot be produced by any other reachable state s j = s i , where two states are said to be reachable if there exists a sequence of actions sampled from the policy that enable the agent to reach state s i from s j with nonzero probability. In other words, each observation uniquely identifies the hidden state, and the POMDP is a block MDP with corresponding Markovian observations.

̸

Next, we consider the case where i = j . We have that for all i, k , ∑ ω Pr( ω | s i )Pr( a k | ω )(1 -Pr( s i | ω )) = 0 . Because each term is nonnegative, this is equivalent to Pr( ω | s i )Pr( a k | ω )(Pr( s i | ω ) -1) = 0 . Because we again have that for all observations there exists an action a k ′ with nonzero probability, this means we can choose k = k ′ to find Pr( ω | s i ) = 0 or Pr( s i | ω ) = 1 for all ω, s i . This means that either the state s i does not produce an observation ω , or the observation ω uniquely determines which state the agent is in. For all ω and i = j , either Pr( ω | s i ) = 0 or Pr( s j | ω ) = 1 , so if we restrict our focus to the set O ( s i ) := { ω ∈ Ω : Pr( ω | s i ) > 0 } , we see that ∑ ω ∈O ( s i ) Pr( ω | s i ) Pr( s j | ω ) = ∑ ω ∈O ( s i ) Pr( ω | s i ) · 1 = 1 . For the remaining ω / ∈ O ( s i ) , Pr( ω | s i ) = 0 .

̸

To recap, the first case tells us that for all ω and i = j , ∑ ω ∈ Ω Pr( ω | s i ) Pr( s j | ω ) = 0 . The second case tells us that for all ω and i = j , ∑ ω ∈ Ω Pr( ω | s i ) Pr( s j | ω ) = 1 . Combining both these cases, we see that ∑ ω Pr( ω | s i ) Pr( s j | ω ) = 1 [ s i = s j ] , which we can write more succinctly as: Φ W = I . We call Φ W the state confusion matrix. For block MDPs, observations cause no confusion about which state the agent is in.

Lastly, by going backwards through the proof, we see that the converse is also true. If the system is a block MDP, then Π S = Φ W Π and so K λ π = K λ ′ π .

̸

̸

̸

F Proof that matching transition-policies is equivalent to an MDP.

Assume that TK = TK ′ . We can reduce this as follows:

Claim: if P is a POMDP such that T Π S = T Φ W Π for any policy, then P is equivalent to an MDP.

Recall that the transition function of the effective MDP model is WT Φ , and the rewards are WR SA . Intuitively, if T and WT Φ are equivalent transition functions, it means they induce the same visitation after any number of steps. Note that due to the premise, the following quantities are equal:

and so on.

Multiplying each successive equation by γ i and summing the results gives:

For any given observation, the POMDP model and the effective MDP model make the same predictions about the expected discounted visitation to future state-action pairs. In other words, they have the same successor representation.

Multiplying both sides on the right by an arbitrary reward function R SA ,

we see that the MC value according to the POMDP model (left hand side) is equal to the TD value under the effective MDP model (right hand side). Both models make the same value predictions for any reward function, therefore the models are equivalent.

We show an example of such a POMDP and its equivalent MDP in Figure 7.

$1

S

S1

14

12

12

12

os

S

12

S2

12

14

12

12

G Memory Augmentation and Norm Details

In this section, we clarify theoretical details in augmenting a POMDP with memory and the λ -discrepancy norm.

G.1 Memory-Augmented POMDP

As referenced in Section 4, here we will explain how to define a memory-augmented POMDP from a base POMDP ( S , A , T, R, Ω , Φ , γ ) . Given a set of memory states M , we will augment the POMDP as follows:

where δ x 0 ( · ) is the discrete point distribution centered at x 0 , and × between two distributions denotes the product distribution. To demonstrate what this means, the notation above is equivalent to seeing T M instead as a function mapping to [0 , 1] and defining T M : S M × A M × S M → [0 , 1] , ( s 0 , m 0 , a 1 , m 1 , s 2 , m 2 ) ↦→ T ( s 2 | s 0 , a 1 ) 1 [ m 1 = m 2 ] , where a memory-augmented action ( a 1 , m 1 ) means to take action a 1 and set the memory state to m 1 .

This augmentation scheme uses the memory states M in three ways: as augmentations of states, actions, and observations. The state augmentation concatenates the environment state S with the agent's internal memory state M . Meanwhile, the action augmentation A M provides the agent with a means of managing its internal memory state. Together, these augmentations allow writing the augmented transition dynamics T M , which are defined so as to preserve the underlying statetransition dynamics T while allowing the agent full control to select its desired next memory state. The observation augmentation Ω M provides the agent with additional context with which to make policy decisions, and the observation function Φ M preserves the original behavior of the observation function Φ while giving the agent complete information about the internal memory state.

We define an augmented policy π M and its tensor counterparts Π M and Π S M over observations and states, respectively, as follows:

where

We extend the Pr( s | ω ) tensors to S M and Ω M as follows:

where

One appealing aspect of the above augmentation process is that it makes the agent's control over its memory function explicit, via 'memory-augmented actions' A M = A × M . However, it is perhaps a bit unusual to combine the agent's policy and memory together. An equivalent and perhaps more intuitive formulation decomposes this action-memory policy into a separate action policy π µ : Ω ×M→ ∆ A and a memory function µ : Ω ×M×A→ ∆ M :

Note that we define memory functions as distributions of next memory states for generality.

There are two ways to view this augmentation procedure. In the first view, outlined above, the agent selects both an action a and a next memory state m ′ , and the environment T M responds by sampling the next state s ′ ∼ T ( ·| s, a ) while deterministically setting the next memory state to m ′ . In the second (equivalent) view, the agent selects only action a , with the agent's memory function behavior folded into the transition dynamics T µ . This view leads to the following POMDP quantities:

where µ S ( ·| s 0 , m 0 , a 0 ) := ∑ ω ∈ Ω Φ( ω | s 0 ) µ ( ·| ω 0 , a 0 , m 0 ) is the effective memory function for states.

The µ -specific policy and weight tensors are defined as follows:

where

To recap, the first augmentation scheme transforms the POMDP P = ( S , A , Ω , T, R SA , Φ , p M 0 , γ ) into P M = ( S M , A M , Ω M , T M , R S M A M , Φ M , p M 0 , γ ) . The second view forms an equivalent POMDP P µ = ( S M , A , Ω M , T µ , R S M A µ , Φ M , γ ) , which folds the agent's memory function µ into the transition dynamics.

Since both of these are valid POMDPs and we already have a general expression for the TD( λ ) value function of a POMDP, we can immediately write down the corresponding value functions for each of these types of augmentation. For P M , we have:

and for P µ , we have:

We provide pseudocode for taking this memory-Cartesian product of a given POMDP in Appendix H.3, Algorithm 3. The pseudocode uses the view that aligns with Equation (10), since that view matches our implementation.

G.2 λ -Discrepancy Norm Weighting

The λ -discrepancy as introduced in Definition 1 contains a weighted norm ∥ C · x ∥ p over the observations ω and actions a 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 L 2 norm, and the p = ∞ or max norm, as described below. For C , we have policy-weighted and occupancy weighted cases. In the policy-weighted case,

we assign the ( ω, a ) entry of C the weight (1 , π ( a | ω )) . In the occupancy-weighted case, we assign the ( ω, a ) entry of C the weight (Pr( ω ) , π ( a | ω )) , where Pr( ω ) is proportional to the discounted observation occupancy.

In Figure 4 and Figure 5, we use the policy-weighted L 2 norm. We use this norm due to issues with performing closed-form gradient descent on the T-maze POMDP. In T-maze, states are aliased in two ways: aliasing between the two hallways (goal up or down), and aliasing between the hallway states. Occupancy weighting puts more weight on discriminating between the hallway states, since the hallway observations appear more frequently for most policies. Uniformly weighting the λ -discrepancy puts more weight on the initial observations, which results in memory functions that place more emphasis on resolving λ -discrepancy due to the starting observations. Closed-form experiments with occupancy weighting yielded less consistent results for T-maze. This seems to only be an issue for closed-form calculation of λ -discrepancy, as we show in Appendix I.6 that occupancy weighting in the sample-based setting does not detract from performance in these environments.

In Figure 3, we use the occupancy-weighted max norm to avoid artifacts due the changing observation function (see text in Section 3.2).

In our sample-based settings in Figures 6, 11, 12, and 14, and Table 2, we use the occupancy-weighted L 2 norm, since it is the simplest choice when using samples as it does not require re-weighting or importance sampling.

Regardless of which norm we choose, we can estimate gradients as in Equation 7 because:

In practice, when optimizing the λ -discrepancy, we omit the outer square root and simply compute the mean squared difference between Q-values, since the square-root function is a monotonic transformation and its argument is non-negative.

H Environments and Experimental Details for Closed-Form Memory Optimization

We now describe the proof-of-concept experiments from Section 4 involving closed-form memory optimization on a range of small-scale classic POMDPs. The environments are: T-maze [Bakker, 2001], the Tiger problem [Cassandra et al., 1994], Paint [Kushmerick et al., 1995], Cheese Maze, Network, Shuttle [Chrisman, 1992], and the partially observable version of the 4 × 3 maze [Parr and Russell, 1995].

We first begin by detailing the small-scale POMDPs we use. Then we describe the algorithm used to calculate closed-form λ -discrepancy gradients, as well as procedures for policy improvement after learning a memory function. Apart from the T-maze, all other POMDPs used in Section 4 were taken from pre-defined POMDP definitions [Cassandra, 2003].

We made one slight modification to the Tiger environment that preserves the original environment behavior but adapts the domain specification to match our formalism such that observations are only a function of state. The original Tiger domain used a hand-coded initial belief distribution that was uniform over the two states L / R , and did not emit an observation until after the first action was selected. Thereafter, the observation function was action-dependent, with state-action pair ( L , listen ) emitting observations left and right with probability 0 . 85 and 0 . 15 respectively, and other actions ( L , * ) emitting uniform observations and returning to the initial belief distribution. Since our agent does not have access to the set of states, it cannot use an initial belief distribution. To achieve the same behavior, we modified the domain by splitting each state L / R into an initial state L 1 / R 1 that always emits an initial observation, and a post-listening state L 2 / R 2 that uses the 0 . 85 / 0 . 15 probabilities. We visualize these changes in Figure 8. This type of modification is always possible for finite POMDPs and does not change the underlying dynamics.

Listen

+10

-100

-1

Right

Left

Right

L

L1

L2

Listen

Listen

12

Left

+10

12

-1

-1

12

Left

-100

12

-1

-1

R

R1

Listen

R2

Listen

Right

-1

Left

Right

Listen

+10

-100

H.1 T-maze Details

We use T-maze with corridor length 5 as an instructive example (see Figure 1). The environment has 18 underlying MDP states: one initial state for each reward configuration (reward is either up or down), five for each corridor, one for each junction, and finally one for each terminal state. 11 There are 5 observations in this environment - one for each of the initial states, a corridor observation shared by all corridor states, a junction observation shared by both junction states, and a terminal observation. The action space is defined by movement in the cardinal directions. If the agent tries to move into a wall, it remains in the current state. From the junction state, the agent receives a reward of +4 for going north, and -0 . 1 for going south in the first reward configuration. The rewards are flipped for the second configuration. The environment has a discount rate of γ = 0 . 9 .

+1

-1

+1

1

+1

00

This environment makes it easy to see the differences between MC and TD approaches to value function estimation. We visualize these differences in Figure 9. MC computes the average value for each observation by averaging the return over all trajectories starting from that observation. By contrast, TD averages over the 1-step observation transition dynamics and rewards, and bootstraps off the value of the next observation. For a policy that always goes directly down the corridor and north at the junction, this leads to an average (undiscounted) return for the blue observation of +4 with MC and (4 -0 . 1) / 2 = 1 . 95 with TD.

H.2 Analytical Memory-Learning and Policy Improvement Algorithm

Algorithm 1 describes our memory optimization procedure, which reduces the λ -discrepancy to learn a memory function, then learns an optimal memory-augmented policy.

Algorithm 1 Memory Optimization with Value Improvement

Input: Randomly initialized policy parameters θ π , where Π = softmax( θ π ) , randomly initialized memory parameters θ µ , POMDP parameters P := ( S , A , T, R SA , Ω , Φ , p 0 , γ ) , number of memory improvement steps n steps ,µ , number of policy iteration steps n steps ,π , learning rate α ∈ [0 , 1] , and number of initial random policies n .

Output: Optimized memory-augmented policy parameters θ π µ and memory parameters θ µ .

// Initialize random policies and select policy to fix for memory learning. { θ 0 , . . . , θ n -1 } ← randomly init n policies ( n ) θ π ← select argmax lambda discrepancy ( θ 0 , . . . , θ n -1 } ) // Repeat policy over all memory states. θ π µ ← repeat ( θ π , |M| ) // Improve memory function. θ µ ← memory improvement ( θ µ , θ π µ , P , n steps ,µ ) // Augment POMDP with memory function. P µ ← expand over memory ( P , θ µ ) // Improve memory-augmented policy over memory-augmented POMDP θ π µ ← policy improvement ( θ π µ , P µ , n steps ,π ) return θ π µ , θ µ

In Algorithm 1, the policy improvement function can be any function which improves a parameterized policy. We use an analytical version of the policy gradient [Sutton et al., 1999] algorithm. The randomly init n policies function returns n randomly initialized policies, and the select argmax lambda discrepancy function picks the one with the largest λ -discrepancy. The memory improvement and expand over memory functions are defined in Appendices H.2.1 and H.3, respectively.

We have noticed that a larger λ -discrepancy tends to provide a better signal for learning memory functions. Sampling random policies for memory improvement is highly likely to reveal a λ -discrepancy, as suggested by the theory in Section 3.2. For this reason, we consider many random policies ( n = 100 ), then use the policy which had maximum λ -discrepancy as the basis for memory optimization.

H.2.1 Memory Improvement Algorithm

In this section we provide pseudocode (Algorithm 2) for the memory-learning algorithm described in Section 4 and used in Algorithm 1. This function takes as input a POMDP P and memory parameters θ µ , and minimizes the λ -discrepancy of Definition 1. This minimization is achieved through a gradient descent update computed using the auto-differentiation package JAX [Bradbury et al., 2018].

Algorithm 2 Memory Improvement

Input: Fixed policy parameters θ π µ , where Π S M µ ( S M × S M × A ) is the expanded effective policy over state for the policy π µ = softmax( θ π µ ) ( S M ×A ) , memory parameters θ µ , POMDP parameters P := ( S , A , T, R SA , Ω , Φ , p 0 , γ ) , number of improvement steps n steps ,µ , learning rate α ∈ [0 , 1] .

Output: Optimized memory parameters θ µ .

Here, update params() is any gradient-descent-like update, such as stochastic gradient descent, Adam, etc. As a note, all parameters θ in these experiments are initialized with a Gaussian distribution, with mean 0 and standard deviation 0 . 5 .

H.3 Memory Cartesian Product

In this section, we define the memory-Cartesian product function, expand over memory() , used by Algorithms 1 and 2. This function computes the Cartesian product of the POMDP P and the memory state space M , as described in Appendix G.1.

Algorithm 3 Memory Cartesian Product ( expand over memory )

Input: Memory parameters θ µ (with corresponding memory function µ ), POMDP parameters SA

P := ( T, R , Φ , p 0 , γ ) , number of memory states |M| // Repeat reward function for each state over each memory m ∈ M ( S M ×A ) . R S M A µ ← repeat over states ( R SA , |M| ) // Calculate transition function memory cross product. // First, calculate the effective memory function over state ( S × A × M × M ) . µ S ← einsum ( ′ ij, jklm → iklm ′ , Φ , µ ) // Now expand the state transition function to include memory state transitions ( S M ×A×S M ) . T µ ← einsum ( ′ iljk, lim → lijmk ′ , µ S , T ) // Calculate observation function memory cross product ( S M × Ω M ) . I |M| is the identity function over |M| . Φ M ← kron (Φ , I |M| ) // Finally, calculate the initial state distribution ( S M ) . p M 0 ← [ p 0 ( s ) if m = 0 else 0 for s, m ∈ S , M ] return P µ = ( S M , A , Ω M , T µ , R S M A µ , Φ M , p M 0 , γ )

Note that einsum is the Einstein summation, and kron is the Kronecker product. The augmented initial state distribution is simply the same distribution as p 0 , except with 0 probability mass over all non-zero memory states, since the memory state always initializes to memory state 0.

H.4 Closed-Form Memory Learning Experiment Details

For all experiments in Section 4, we run memory optimization on the suite of POMDPs with the following hyperparameters. We optimize memory for n steps ,µ = 20 K steps and run policy iteration for n steps ,π = 10 K steps. For all gradient-based experiments, we use the Adam optimizer [Kingma and Ba, 2015].

For the belief-state baselines, solutions were calculated using a POMDP solver from the pomdp-solve package [Cassandra, 2003]. The performance of the belief-state optimal policy was calculated by taking the dot product between the initial belief state and the maximal alpha vector for that belief state. This returns a metric comparable to the initial state distribution weighted value function norm, which we use as a performance metric for our memory-augmented agents.

The belief-state solution for the 4 × 3 maze was solved using an epsilon parameter of ϵ = 0 . 01 , due to convergence issues with the environment when utilizing the POMDP solver.

H.5 Closed-Form Experiments on Parity Check

In Section 3.3 we introduced the Parity Check environment as an example of a POMDP with zero λ -discrepancy for all policies. There we showed that in such environments, we can still detect partial observability by using randomly initialized memory functions. But the question remains: how prevalent are such environments?

We find that these sorts of environments appear to be quite rare. The symmetry disappears if we modify the example even slightly, such as by changing the start state probabilities or by introducing a 'stay-in-place' action at any one color. These modifications lead to the λ -discrepancies in Figure 10, which offer some reassurance that such edge cases are the exception and not the rule.

Besides modifications to the POMDP itself, adding memory can induce a non-zero λ -discrepancy. In particular, random memory functions reveal a λ -discrepancy in the system that can be minimized to learn memory, as demonstrated in Figure 4. We now describe this memory optimization procedure.

Memory Learning Memory optimization for the Parity Check experiments in Figure 4 (right) followed almost the same procedure as Algorithm 1 except for selecting the initial policy. In this case, any randomly initialized memoryless policy has zero λ -discrepancy, but augmenting with a randomly initialized memory function reveals a non-zero measure. Thus, for this experiment, we augment our POMDP with a random memory function before choosing the initial policy that maximizes the λ -discrepancy. All other algorithmic details remained the same.

个

+1

个个

K

100

100

discrepancy

-discrepancy

个个

→

10-6

10-6

→

10-12

10-12

-1

Pr(sTAY|BLUE)

Po(RED → PINK)

↑个

10-18

10-18

0.0

0.5

1.0

0.0

0.5

1.0

→

+1

Probability

Probability

I Scaling the λ -Discrepancy with PPO and RNNs

I.1 Algorithm Details

We build our λ -discrepancy minimizing reinforcement learning algorithm on top of online recurrent PPO. Normally, PPO would have two losses: one for the actor and one for the critic.

Here, L CLIP is the clipped surrogate loss for the actor, defined by:

where ρ t ( θ ) is the ratio between policy probabilities of the old policy and new policy, ˆ A t is the estimated advantage at timestep t , and c V and L V are the value coefficient and value function loss, respectively. L V is the mean-squared TD error, with targets calculated using a truncated version of λ -return [Sutton and Barto, 2018]. Finally, c Ent and H [ π µ ( · | ω t , z t -1 )] are the entropy coefficient and entropy of the policy π µ that is parameterized by θ RNN and θ Actor . This term is used for exploration through entropy maximization [Mnih et al., 2016]. For our memoryless baseline, we use a nonrecurrent architecture that is a function of only the current observation (and optionally the most recent action).

Our proposed algorithm replaces L V with L V,λ 1 ,λ 2 , which has three components: two value losses for two separate value heads (each estimating λ -returns for different λ s), and our λ -discrepancy auxiliary loss.

where L Λ ( θ ) is our λ -discrepancy loss, as defined in Equation 8, and L V,λ is the λ -return-as-target TD error for a given λ . Here, β ∈ [0 , 1] is a hyperparameter trading off between the λ -discrepancy loss and the value losses. Finally, the agent only uses V λ 1 as its value estimate for its advantage function.

The computational and memory requirements for PPO with and without the λ -discrepancy auxiliary loss are similar. While the λ -discrepancy loss requires training a second value head, this only adds two additional layers to the critic network and results in a small number of additional parameters (about 12% overhead). As for computation, the addition of this auxiliary loss only adds a small constant factor to the backpropagation of gradients, since we learn everything end-to-end. Overall, the wall-clock time of runs between the baseline algorithm and our algorithm is comparable.

I.2 Environment Details

All four of the environments used for evaluation in Section 5 were all re-implementations of environments used in Silver and Veness [2010], in JAX [Bradbury et al., 2018], allowing for massive hardware acceleration. We now give details of these environments.

I.2.1 Battleship

Partially observable Battleship, a one-player, limited observation variant of the popular two-player board game, was first introduced as a benchmark for partially observable Monte-Carlo planning [Silver and Veness, 2010]. This variant is particularly challenging since the agent has to reason about the unknown positions of the ship and keep track of past shots. The observation space is Ω = { 0 , 1 } ; at every step, the agent receives a binary observation: 0 if the last shot missed a ship, and 1 if the last shot hit a ship. The state space in this game contains all possible board states, which is astronomically large. The action space is A = { 1 , . . . , 10 } × { 1 , . . . , 10 } , or all possible grid locations. The agent is only allowed to take valid actions at every step (no position can be selected twice), which is achieved through action masking of the actor. This makes the problem a finite horizon problem, with a horizon at most 10 × 10 = 100 . The environment terminates when all positions on the grid with a ship are hit. The rewards at every step are -1, with a positive reward of 10 × 10 = 100 when all ships are hit/the environment terminates. The discount factor γ is set to 1 here, since we are in the finite horizon setting.

At every environment reset, 4 ships of length (5 , 4 , 3 , 2) are uniformly randomly placed on a 10 × 10 grid. At every step, to allow for easier learning of all agents, we concatenate the last action to the observation, so that our RNN conditions on the entire history, as opposed to only observations.

I.2.2 Partially Observable PacMan

Partially observable PacMan (a.k.a. PocMan ), a variant of the popular arcade game, was also introduced to test partially observable Monte-Carlo planning with a simulator [Silver and Veness, 2010]. In this version, the agent can only observe indicators of its surroundings, including walls in its cardinal directions, whether there is a ghost in its line of sight, a power pellet nearby, or food nearby. These highly obscured agent-centric observations require the agent to localize within the map, seek food and power pellets, and avoid ghosts with only local sensor information.

Concretely, the observation space is an 11-dimensional binary vector. The first four values are 1 if there is a wall in each of the four cardinal directions of the agent. The next element is 1 if there is food Manhattan distance 2 or less away from the agent. The next value is 1 if there is a ghost a Manhattan distance 2 or less away. The next four elements are 1 if there is at least one ghost in the line of sight in each of the four cardinal directions. The last element is if the agent currently has a power pellet.

As with Battleship, the state space in PocMan is also very large. The map is a 19 × 21 sized map defined in the original instantiation of the environment [Silver and Veness, 2010], and the state consists of the agent and ghost locations, as well as the binary status of every dot and power pill. The action space has four actions, corresponding to moving in the four cardinal directions. The agent receives a reward of +200 if it eats a ghost, +10 if it collects a pellet, and +20 if it collects a power pellet. The discount factor is set to γ = 0 . 95 .

The episode terminates when the agent is killed by a ghost, the agent collects all pellets, or the timer for the environment elapses. At the beginning of each episode, all ghost locations are reset, and the agent is placed in the same fixed starting position. In this setting, we also append the agent's most recent action to its observation.

In our work, we implement this domain by adding the observation function on top of the PacMan environment in the Jumanji reinforcement learning framework [Bonnet et al., 2024].

I.2.3 RockSample

RockSample [Smith and Simmons, 2004], a well-known partially observable benchmark in POMDP planning literature, is a rock-collecting problem simulating a rover scanning potential rock samples and collecting them if they are desirable. RockSample ( n, k ) is a problem with a grid of size n × n , and k rocks distributed randomly throughout the environment. The two variants we test on have a grid size of 11 × 11 and 15 × 15 , as well as 11 and 15 rocks in each environment, respectively. When instantiated, each RockSample environment for each seed randomly samples rock positions. At the reset of each environment, each of these rocks are uniformly randomly assigned to be either a good or bad rock. The action space in this environment is 4 + 1 + k = 5 + k . The first 4 actions correspond to moving in the cardinal directions. Action 5 corresponds to sampling in the current position. The last k actions correspond to checking each of the k rocks. Checking a rock will probabilistically tell the agent the correct parity of the rock, depending on the half-efficiency distance d hed and the l 2 distance d i ( s ) between the agent's position and the rock being checked, given the current state s :

In this work, we simply set d hed to be the maximum possible distance between any two points in the grid. Traversing to and then sampling a good rock gives a reward of +10. Sampling a bad rock gives a reward of -10. Exiting to the east border of the environment will terminate the environment and result in a reward of +10. After sampling a good rock, sampling the same rock again will result in a reward of -10. This environment has a discount rate of γ = 0 . 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 a vector of size 2 n + k binary values, where the first 2 n elements are a two-hot representation of the xy -coordinates of the agent. The final k observations are set after choosing a check action, and set a 1 at the position of the checked rock if it appears to be good, and a 0 if it appears to be bad. This observation function is a slight departure from the previous RockSample problem definition [Smith and Simmons, 2004, Silver and Veness, 2010] in two aspects, based on an implementation of RockSample that is used to test RNNs [Tao et al., 2023], as opposed to model-based planning algorithms. The first is that after an agent checks a rock and the agent receives a positive sensor reading, the observation bit corresponding

to this rock is set to 1, and remains 1 until sampled. This makes the memory learning portion of the problem easier, but the agent is still required to remember which rocks it has sampled before, and also deal with the stochasticity of checking. The second change is to not have an explicit negative sensor reading. Instead, after checking a rock, the observation bit corresponding to the sensor reading of the rock remains 0. This makes the problem harder, since the agent has to infer the negative sensor reading, and disambiguate it from a null sensor reading, when the agent takes non-check actions.

I.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 is finite-horizon and undiscounted.

RockSample 11 x 11

RockSample 15 x 15

Return

20

Discounted

10

10

0

0

0

1

2

3

0.0

0.2

0.4

0.6

0.8

1.0

Timestep

×106

Timestep

×107

Battleship10× 10

Partially Observable PacMan

Discounted Return

300

15

200

PPO+RNN+LD

10

PPO+RNN

100

PPO

5

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

X107

×107

Timestep

Timestep

RockSample 11 x 11

RockSample 15 × 15

Episode Reward

20

20

0

0

0

1

2

3

0.0

0.2

0.4

0.6

0.8

1.0

Timestep

X106

Timestep

×107

Battleship 10 × 10

Partially Observable PacMan

Episode Reward

15

2000

10

PPO+RNN+LD

1000

PPO+RNN

PPO

5

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

Timestep

×107

Timestep

×107

I.4 Experimental and Hyperparameter Details

Our base PPO algorithm is an online PPO algorithm that trains over a vectorized environment, all parallelized using JAX [Bradbury et al., 2018] and the PureJaxRL batch experimentation library [Lu et al., 2022] with hardware acceleration. The hyperparameter sweep was performed on a cluster of NVIDIA 3090 GPUs, and the best seeds presented were run on one GPU for each algorithm, running for 1 to 12 hours, depending on the domain. Our environment is vectorized over 4 copies, and we use truncated-backpropogation through time [Jaeger, 2002] as our gradient descent algorithm, with a truncation length of 128 . Our L CLIP clipping ϵ is set to 0 . 2 . The value loss coefficient is set to c V = 0 . 5 . We also anneal our learning rate over all training steps, and clip gradients when the norm is larger than 0.5 by their global norm [Pascanu et al., 2013].

We now detail the hyperparameters swept across all environments and all algorithms. We do so in Table 1.

We conduct this sweep across all environments for 5 seeds each, and use the highest area under the learning curve (AUC) score in each environment to select the best hyperparameters for each environment. Note, we swept over 10 seeds for recurrent PPO and λ -discrepancy-augmented recurrent PPO for the PocMan environment, due to high variance in returns. After hyperparameter selection, we re-run all algorithms on all environments with the best selected hyperparameters for 30 different seeds to produce the learning curves in Figure 6. We present the best hyperparameters found for both algorithms in Figure 2.

(a) Best hyperparameters found for λ -discrepancy-augmented recurrent PPO.

(b) Best hyperparameters found for the recurrent PPO baseline.

(c) Best hyperparameters found for the memoryless PPO baseline.

Table 2: Best hyperparameters for each environment and each algorithm. Hyperparameters were found using 5 seeds, and taking the maximum AUC.

Our network architectures are standard multi-layer perceptions (MLPs). Both actor and critic networks are two-layer MLPs with ReLU [Nair and Hinton, 2010] activations between layers. The actor network applies a softmax function to its output logits. Our recurrent neural network is a dense layer, followed by ReLU activation, then the GRU cell, then another dense layer. For Battleship, to better condition on the hit-or-miss observation, we add an additional dense layer after the first dense layer of the RNN (which has 2 × latent size for this environment), that takes as input the outputs of the previous layer

concatenated with the hit-or-miss bit (a residual connection). For our memoryless baseline, we replace the recurrent neural network with a 3-layer MLP with ReLU activations between each layer. Put together, our actor-critic network takes in an input observation (and optionally, the previous action) and passes it, as well as the previous latent state, through the GRU for a new latent state. For the memoryless baseline, the 3-layer MLP simply encodes the observation (and potentially action) into the latent state, with no recurrence. It uses this new latent state as inputs to both the actor and critic.

We now detail environment-specific hyperparameters in Table 3. 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/brownirl/lambda_discrepancy .

I.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 to collect both RNN + PPO hidden states and the RNN + PPO + LD hidden states. With this dataset, the each probe was trained with the corresponding RNN hidden states as inputs, with the pellet occupancy of all potential pellet positions as the target for the prediction. Both probes are 3-layer neural networks with ReLU activations [Nair and Hinton, 2010] between each layer, and a sigmoid function over the final logits to map outputs to 0 and 1. We use a binary cross entropy loss between these predictions and the pellet occupancy targets. The hidden size of the network was 1024, with a step size of 0.0001. At every step of training, a batch of 32 time steps are uniformly randomly sampled from the dataset, and the binary cross entropy loss was minimized. The agent performs 10M steps of training to reach the performance visualized.

I.6 Small-Scale POMDP Experiments

We also run both recurrent PPO and our λ -discrepancy-augmented PPO algorithm on the small-scale POMDPs evaluated in the analytical experiments in Section 4. For these experiments, all agents had a latent state size of 32 , with c Ent = 0 . 05 , and the same hyperparameters swept as in Appendix I.4. We show results in Figure 14.

These results imply that the baseline algorithm is already sufficient for solving these tasks, and any additional auxiliary losses cannot help with performance, since performance is already near optimal. We would like to note that, as with the results in Figure 6, adding our λ -discrepancy auxiliary loss never harms the performance of PPO.

4x3

Cheese

Hallway

0.5

0.75

0.50

0.0

0.50

0.25

-0.5

0.25

0.00

0.0

0.5

1.0

0.0

0.5

1.0

0

1

2

×106

×106

×106

Heavenhell

Network

Paint

Return

1.0

200

0.5

0.5

R

01

0.01

0.0

count

-200

0

1

2

3

0.0

0.5

1.0

0.0

0.5

1.0

S

×106

×106

×106

Shuttle

Tiger

T-maze

4

0

20

-20

2

PPO+RNN+LD

-40

PPO+RNN

0

2

3

0.0

0.5

1.0

0

1

2

×106

×106

×106

Timestep

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 as 'the λ -discrepancy is exactly zero for all Markov decision processes and almost always non-zero for a broad class of partially observable environments' are presented with linked proofs in Section 3. Our analytical results that show mitigation of partial observability are presented in Section 4. Our RNN-based results in Section 5 show that minimizing the λ -discrepancy between two learned value functions as an auxiliary loss alongside traditional reinforcement learning, is often significantly more effective (and never worse) than training the same agent with just a single value function.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

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 are encouraged to create a separate 'Limitations' section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 2 in Appendix E. We also prove one additional claim in Appendix F. We explain the memory augmentation process in Section G.1.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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 or 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 conclusions of the paper (regardless of whether the code and data are provided or 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 5 are described in Appendix I.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 be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a 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, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it 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 via detailed 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.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the 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 a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: As stated in Section 1, our code is available at https://github.com/ brownirl/lambda_discrepancy . This repository both has detailed instructions on how to conduct the sweeps, as well as the best runs over the 30 seeds mentioned in both experiments. All experiments were done by setting seeds for the pseudo-random number generator, so results are deterministic with respect to the seed and exactly reproducible. Data from our runs are also provided in the repo.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a 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 to 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, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: These are present in Appendix H and Appendix I. Specifically, in I.4.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Throughout the work, all error bars are included for 95% confidence intervals over 30 runs, as described in Sections 4 and 5.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer 'Yes' if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Details of compute usage, hardware and time taken for experiments are discussed in Appendix I.4.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We followed the Code of Ethics in its entirety. No human research subjects or participants were utilized. We used no datasets that contained personally identifiable information, required participant consent, or were deprecated. We violated no copyright or fair use licenses, and representative evaluation was not a concern.

Our work is of a fundamental nature that has no specific negative impact for safety as compared to other fundamental reinforcement learning works. It contains no particular threats towards security, discrimination, surveillance, deception and harassment, the environment, human rights, or bias and fairness.

The POMDP environments used and the models using them are described in Section 4 and Section 5. There is no sensitive data requiring strong security measures. Standard

security measures are in place as provided through GitHub. Our models do not have a high risk for misuse. Information required to understand research artifacts and reproduce all results in this paper are provided at the code repository. Legal compliance has been ensured. An Apache License v2.0 is provided for the code at https://github.com/brownirl/ lambda_discrepancy .

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the broader impacts of this paper in Appendix B. Our paper, as a work on the foundations of reinforcement learning, is not tied to any particular applications.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our work focuses on the foundations of reinforcement learning and does not entail a high risk for misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Open-source repositories used in this work have been cited, particularly in Appendix I.4. In addition, our code is open-sourced with the Apache 2.0 license.

Guidelines:

The answer NA means that the paper does not 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 a URL.

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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are 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 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 is well documented, with proper READMEs on experimentation and data parsing.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of 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 is 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:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

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 an 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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\nPlease see weaknesses above.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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?\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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?\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe authors address the limitations of their work in a separate section in the Appendix.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-19"}