The Wasserstein Believer: Learning Belief Updates for Partially Observable Environments through Reliable Latent Space Models
Authors: Raphaël Avalos
Florent Delgrange
Ann Nowe
Guillermo Perez
Diederik M Roijers
Venue: ICLR 2024 poster (2024)
Keywords: pomdp
guarantees
representation learning
reinforcement learning
Primary area: reinforcement learning
TL;DR: Wasserstein Belief Updater is an RNN free RL algorithm for POMDPs that learns a representation of the history via an approximation of the belief update in a reliable latent space model, providing theoretical guarantees for learning the optimal value.

Abstract: Partially Observable Markov Decision Processes (POMDPs) are used to model environments where the state cannot be perceived, necessitating reasoning based on past observations and actions. However, remembering the full history is generally intractable due to the exponential growth in the history space. Maintaining a probability distribution that models the belief over the current state can be used as a sufficient statistic of the history, but its computation requires access to the model of the environment and is often intractable. While SOTA algorithms use Recurrent Neural Networks to compress the observation-action history aiming to learn a sufficient statistic, they lack guarantees of success and can lead to sub-optimal policies. To overcome this, we propose the Wasserstein Belief Updater, an RL algorithm that learns a latent model of the POMDP and an approximation of the belief update under the assumption that the state is observable during training. Our approach comes with theoretical guarantees on the quality of our approximation ensuring that our latent beliefs allow for learning the optimal value function.

Published as a conference paper at ICLR 2024
THE WASSERSTEIN BELIEVER LEARNING BELIEF UPDATES FOR PARTIALLY OBSERVABLE ENVIRONMENTS
THROUGH RELIABLE LATENT SPACE MODELS
Raphael Avalos * Florent Delgrange- 2.2*
Ann Nowé"! Guillermo A. Pérez2.31 Diederik M. Roijers
1 Al Lab, Vrije Universiteit Brussel (Belgium) 2 University of Antwerp (Belgium)
3 Flanders Make (Belgium) City of Amsterdam (The Netherlands)
(raphael.. avalos, florent. delgrange @
ABSTRACT
Partially Observable Markov Decision Processes (POMDPs) are used to model environments where the state cannot be perceived, necessitating reasoning based past observations and actions. However, remembering the full history is gener- ally intractable due to the exponential growth in the history space. Maintaining a probability distribution that models the belief over the current state be used a sufficient statistic of the history, but its computation requires access to the model of the environment and is often intractable. While SOTA algorithms use Recur- rent Neural Networks to compress the observation-action history aiming to learn sufficient statistic, they lack guarantees of success and can lead to sub-optimal policies. To overcome this, we propose the Wasserstein Belief Updater, an RL algorithm that learns latent model of the POMDP and an approximation of the belief update under the assumption that the state is observable during training. Our approach comes with theoretical guarantees on the quality of our approximation ensuring that our latent beliefs allow for learning the optimal value function.
INTRODUCTION
Partially Observable Markov Decision Processes (POMDPs) define a powerful framework for mod- eling decision-making in uncertain environments where the state is not fully observable. These problems are common in many real-world applications, such as robotics (Lauri et al.] 2023). and recommendation systems (Wu al. 202I). In contrast to Markov Decision Processes (MDPs), in POMDP the agent perceives an imperfect observation of the state that does not suffice as condition- ing signal for an optimal policy. As such, optimal policies must take the entire interaction history into account. As the space of possible histories scales exponentially in the length of the episode, using histories condition policies is generally intractable. An alternative the notion of belief, which is defined sufficient statistic of the history (Kaelbling et al.) 1998) but the computation of their closed-form expression require the access to model of the environment and is in general intractable.
probability distribution over states based on the agent history. Beliefs are
To overcome those challenges, SOTA algorithms compress the history into a fixed-size vector with the help of Recurrent Neural Networks (RNNs) (Hausknecht & Stone, 2015). Nonetheless, this may lead to information loss, resulting in suboptimal policies. To improve the likelihood of obtaining sufficient statistic, RNNs can be combined with regularization techniques, including generative models (Chen et al., 2022] Hafner et al. 2019 (021), particle filtering (Igl et al.) 2018 Ma et al.) 2020), and predicting distant observations (Gregor al. 2018 2019h. However, none of these tech- niques guarantee that the representation of hhistoriees induced by RNNs is suitable for optimizing the return. Additionally, many algorithms assume that beliefs are simple distributions (c.g.. Gaussian),
which limits their applicability (Gregor et al. [2018] Lee et al. 2020 Hafner et al. 2021).
In this paper, we propose Wasserstein Belief Updater (WBU), model-based reinforcement learning (RL) algorithm for POMDPs that allows learning the belief space over the unobservable states. Specifically, WBU learns an approximation of the belief update rule through latent space model
Both authors contributed equally this research, alphahetic order.
Equal supervision.

---

Published as a conference paper at ICLR 2024
whose behaviors (expressed as expected returns) close to those of the original environment. Furthermore, we show that WBU is guaranteed to induce a suitable representation of the history to optimize the return. WBU consists of three components that are learned in round-robin fashion: the model, the belief learner, and the policy (Fig. D. Harnessing only histories to learn a model whose dynamics can be provably linked to the original unobservable environment poses considerable challenge. Therefore, in the same spirit as the Centralized Training with Decentralized Execution paradigm in multi-agent RL (MARL) (Oliehoek al. 2008) Avalos et al. 2022), where leveraging additional information such as the true state of the environment is common practice, we assume that the POMDP states can be accessed during training. While this might seem restrictive at first sight, this assumption is typically met in simulation-ball -aaed training and can be applied in real-world settings such as robotics, where extra sensors can be used during training in a laboratory setting.
Our core contribution is the development of a sound framework equipped with theoretical guarantees in the context of RL within partial observability. While SOTA algorithms primarily concentrate on enhancing the overall return- potentially resulting in substantial performance gains we contend that performance is not the exclusive goal and that possessing guarantees is equally important, as the balance between these two aspects varies based on the specific application. By tackling POMDPs with a formal approach, we offer theoretical guarantees that other methods cannot provide: we ensure that our latent model is able to replicate the dynamics of the original, partially observable environment, which further yields a belief representation suitable for learning the value function.
We learn the latent model of the POMDP via a Wasserstein auto-encoded MDP (WAE-MDP, Del- grange et al. 2022), which embeds bisimulation metrics intuitively leading to our guarantees. In parallel, we maintain a belief distribution over the latent state space via a belief update network: we minimize its Wasserstein distance to the exact belief update rule, through a tractable variational proxy. To allow for complex belief distributions, we use normalizing flows (Kobyzev et al. [02I]. In contrast to SOTA algorithms, the beliefs are only optimized towards accurately replicating its up- date rule. While we call recursively the belief network to maintain the belief distribution, we do not back-pree pock-propagate through time and thus implement it as a simple feed forward network. The policy is learned on the latent belief space using a vector integrating the parameters of the belicf distribution. Our experimental results are promising and show the ability of our algorithm to learn to encode the history into a representation useful to learn a policy, without using RNNs
Other related work. DVRL (gg et al. 2018) extends A2C (Mnih et al. 2016) combined with RNNs (R-A2C) with auxiliary losses aiming to learn beliefs via a variational autoencoder and particle fil- tering, but lacks guarantees. DVRL further assumes independent normal distributions for beliefs, limiting its applicability. FORBES (Chen et al. 2022) use normalizing flows but learn policies con- ditioned on latent states, which suboptimal as the state distribution is approximated with a single sample. Some works focus on specific POMDP types. like compact image representations (c.g., visual motor tasks, Lee et al. (2200) or states masked with Gaussian noise (Wang & Tan [021]. While accessing states is common in MARL, it is less common in single-agent but has been explored in kernel-POMDPs (Nishiyama et al., 2012), and more recently for sample efficient learning (Lee)
al., 2023). Leveraging additional in information available during training (not necessarily the states) has also been explored by [Lambrchhts et al. (0023), but RNNs remain crucial while no abstraction nor representation quality guarantee
provided. Finally, other works (Gelada et al. [019 Del- grange et al.] (202)) study similar value difference bounds to ours and connect them to bisimulation theory (Larsen & Skou 1989; Givan al. 2003), but in fully observable environments.
BACKGROUND
2.1 PROBABILITY DISTRIBUTIONS AND DISCREPANCY MEASURES
We write (x) for the set of all measurable subsets of complete, separable space X, A(X') for the set of measures on X, and d, E ((X)) for the Dirac measure with impulse a E X. Let P, Q E A(X) be measures with densities p, q. the divergence between P andd Q can be measured via:
the solution of the optimal transport problem (OT), defined as Wa (P.Q) = inf Ez,y~A c(2, y), which is the minimum cost of changing P into Q (Villani 2009), where c: X x X - [0, 20) is cost function and the infimum is taken over the set of all couplings of P and Q. When c is equal to distance metric d over X, Wa is the Wasserstein distance between the two distributions.
2

---

Published as a conference paper at ICLR 2024
mosked outoropresiiee flens
nest
R
foo gondiens (lnd)
(8,0)
3,00
Bept
M
-10
behief updute
7(- by Me, 0r-1)
Ee P(/fi (Baaaec) Dlon-a13 Pr Poleee Cot 188)
R
P
Wasserdein Belief Updater
f
pamunes of Kcl
Replay Buffer
WAE-MDP
A2C Agent
Environment
D
(s.0,a.r.0f.0
π
P
M
Figure WBU framework. The WAE-MDP is presented in Sect. 3 and WBU in Sect. Learn- ing the different components is done in round-robin fashion. The WAE-MDP learns from data collected by the RL agent and stored in Replay Buffer. WBU uses the transition function P and observation decoder O of the WAE-MDP learn to approximate the belief update rule. The agent learns policy conditioned on the resulting sub-belief Be (i.e., the parameters of the latent belief b,).
the total variation distance (TV), defined as drv (P,Q) = sup PP(A) Q(A)) If is equipped with the discrete metric 1 TV coincides with the Wasserstein metric. the Kallback-Leibber (KL) divergence, defined as DKL Q) Ee- [log(p(2\/((=)]
2.2 DECISION MAKING UNDER UNCERTAINTY
Markov Decision Processes (MDPs) are tuples M = AS P. 81,7) where S is a set of states; A, set of actions; P: A (SS), a probability transition function that maps the current state and action to distribution over the next states; R: S A R, reward function; initial state; and 7 E [0, 1) discount factor. spaces as continuous MDPs. In that case, we assume S and A are complete separable metric spaces equipped with Borel o-llgerra. An agent interacting in M produces trajectories. i.c., sequences of states and actions (Su.r; 0g.r ,> where Su 8s and $1+1 P(- St, a4) for t < T.
E S, the refer to MDPs with continuous state or action
Policies and probability measure. stationary policy म: S A(A) prescribes which action to choose at each step of the interaction. A policy æ and M induce unique probability measure TM on the Borel a alaleebra over (measurable) infinite trajectories (Puterman, 1994). The typical goal of an RL agent is to learn policy that maximizes the expected return, given by Ex [Ei X R(s,, a.)] by interacting with M. We omit the superscript when the context clear.
Partially Observable MDPs (POMDPs) are tuples P (M, 22 O) where M is an MDP with state space S and action space A: Ω is set of observations; and O: S x A - A(2) is observation function that defines the distribution observations that may occur when the MDP M transitions to state upon the execution of particular action. An agent in P actually interacts in M, but without directly observing the states of M: instead, the agent perceives observations, which yields histories, i.e., sequences of actions and observations (0g:7-1, 0,.r) that can be associated an (unobservable) trajectory (Sorr, Oo-r-1) in M, where ot+1 O(- 8t+1, at) for all / T.
Beliefs. Unlike in MDPs, where it is sufficient to take decisions based on the current state (Puterman, 1994), policies solely based on the current observation P are not sufficient to optimize the return. Intuitively, due the partial observability, the agent must take into account full histories in order
3

---

Published as a conference paper at ICLR 2024
to make an informed decision its next action. Alternatively, the agent maintain belief by E A(S) = B over the current state of M (Aström 1965). Given the next observation 0c+1. the
next belief 1 is computed according to the belief update function B x A x Ω B, where T The,, az, 0++1) = b iff the belief over any next state se+1 E S has for density
(1+1() Ba,nb, P(1 Sc, oz) O St+1.0e)
(1)
Eer-br Pg O(0 s',ae)
Each belief 1 constructed this way is sufficient statistic for the history (au..) Orieat) to optimize
the return (Kaelbling et al.) 1998). We write -$ T*(ae: 01.++1) ) = T(: (t, oc+1) 1e+1)a o.... 7 4D 01) b for the recursive application of along the history. The belief update rule derived from allows to formulate P as continuous belief MDP MM = AB 8. Rg, b1,7), where Ps(b b, = E (-s.s-()))s.a) (or-0((ss,a)) or().a,o') (b'); R2(b,a)
Es- R(s, a); and by = d2y As for all MDPs, MM and any stationary policy for MM (thus conditioned on beliefs) induce well-defined probability space over trajectories of Ma. which allows optimizing the expected return in P.
2.3 LATENT SPACE MODELING
Latent MDPs. Given the original (continuous or very large, possibly unknown) environment M, latent space model is another (tractable, explicit) MDP M = (S.A,P. P, R, 81,7) with state space linked to the original one via a state embedding function: o: S
Wasserstein Auto-encoded MDPs (WAE-MDPs, Delgrange et al. 2023 are latent space models that are trained based on the OT from trajectories resulting from the execution of the RL agent policy in the real environment M, to that reconstructed from the latent model M.. The optimization process relies on a temperature E (0, 1) that controls the continuity of the latent space learned, the zero- temperature limit corresponding to discrete latent state space (see Appendix [ discussion). This procedure guarantees M to be probably approximately bisimilarly close (Larsen & Skou 1989. Givan et al., 2003 Delgrange et al., 2022) to M as X - O: in nutshell, bisimulation metrics imply the closeness of the two models in terms of probability measures and expected return (Desharnais
et al. 2004 Ferns al. 2011). Specifically, a WAE-MDP learns the following components:
state embedding function a latent reward function
8:S-
latent transition function
P: A
A(5)
R:SxAL
a state decoder
v: S S S.
(2)
LEARNING THE DYNAMICS
The agent is assumed operate within POMDP. In an RL setting, the former have no explicit access to the environment dynamics: instead, it reinforces its behaviors through interactions and ex- periences without directly accessing the transition, reward, and observation functions of the environ- ment. To provide the aforementioned guarantees, we henceforth adhere the following assumption.
Assumption (Access to the state during training). In addition to the observation, the agent is able to observe the true state of the environment, but only during the training phase.
Remark 1. Seemingly restrictive, Assumption | can actually be met in a broad range of training scenarios, in particular those relying on simulators, where one could merely consider the RAM as the state. Otherwise, additional sensors with higher fidelity could be considered to obtain the state. Other applicable scenarios are model-based design or model prediccve control, where model is accessible during training, and situations where accessing the state is costly (Bulychev et al. 2012).
Concretely, when the RL agent interacts in POMDP P = (M,Ω,)) O) with underlying MDP M = (S, R. 81, Y), we leverage this access to allow the agent to learn the dynamics of the environment, i.e., those of M, as well as those related to the observation function O. To do so, we learn an internal, explicit representation of the experiences gathered, through a latent space model. We then this model as a teacher for the agent make it learn how perform its belief updates. Hence, acquiring an accurate environment model is crucial to learn a reliable belief update function. In Sect. 3.2 further demonstrate that the resulting model is guaranteed to closely replicate the original environment behaviors. The trick we use to learn such a model is to reason on an equivalent POMDP, where the underlying MDP is refined to encode all the crucial dynamics.
4

---

Published as a conference paper at ICLR 2024
3.1 THE LATENT POMDP ENCODING
We enable learning the dynamics of P via a WAE-MDP by considering the POMDP P (Ma, Ωt O³), where (i) the state space of the underlying MDP is refined to encode the observa- tions: Ma = (Sn,A, Pg. R2. ($1, 01),7) with Sp = S x Ω. Pa(s', o' 8, o, 0, P( 8, 8,a). O(d') s',, a). Rs( 0) a) R(s a). and or is an observation from Q linked to the initial state sy; (ïi) the observation function O1 Sq - 52 is the deterministic projection of the refined state on the observation space, with ('((s, o)) o. The POMDPs P and P are equivalent (Chatterjee et al,) 2016]: Pt captures the stochasticity of O in its transition function through the refinement of the state space, further yielding a deterministic observation function, only dependent on refined states.
Henceforth, we reduce the problem of learning a latent space model of P to learning a WAE-MDP from Man Precisely, we learn a latent MDP M embedding Φ: Sa S. The latent MDP M encodes the observation dynamics through P and enables to learn the deterministic observation function O' through the decoder 3:: S - Sp. by decomposing the latter in networks v/$: S S and Oμ: S Ω, further yielding (s)) (5), (5)) for all E S. This way, the WAE-MDP learns all the components of Pt, the latter being equivalent to P. With this model, we construct latent POMDP P OM where O is a Dirac measure with impulse Tn. Figure 6 in Appendix presents a visualization of the relationship between the different models.
5177) linked to Ma via the
As with any POMDP, the belief update function 7 of P allows reason on the belief space to optimize the return. The belief update procedure is illustrated in Appendix B Formally, assuming the latent belief at time step t 3 0 is by E A(S) B, a is executed, and then o--1 observed, 6, is updated according to F(be, as, 0c+1) ) = b iff, for any (unobservable) next state $ E S,
(++1(8++1) Ba,nh P(5 5t, at) (o0+ $2+1)
(3)
Bs,-h, (g O 8')
Latent policies. Given any history h (A- (A) running latent policy म: B ((A) in
is possible by converting h into latent belief (*hh) b and executing the action prescribed by #(- 6). Training M grants access to the dynamics required to update the belief through its closed form 3... However, integrating over the full latent space remains computationally intractable.
As a solution, we propose to leverage the access to the dynamics of M to learn a latent belief encoder :p: B x A x Ω B that approximates the belief update function by minimizing D(*((h) (h), (h)) for some discrepancy D and E (A 22)* drawn from some distribution. The belief encoder ip thus enables to learn policy म conditioned on latent beliefs to optimize the return in P: given the current history h, the next action to play is given by a ~ #( 10" (
Two key questions arise: (i) does the WAE-MDP encoding induces a latent POMDP with behaviors close to P? (ii) the history representation induced by & suitable for optimizing the expected return in P? Our guarantees hinge on the history distribution and chosen discrepancy. The next section details our theoretical analysis of the required distributions and losses to answer (i) and (ii).
3.2 LOSSES AND THEORETICAL GUARANTEES
To yield the guarantees, specifically target the episodic RL process setting for drawing histories. Assumption (Episodic RL process). The environment P embeds a special reset state so that (i) under any policy, the environment is almost surely eveniually reset; (ii) when reset, the environment transitions to the initial state; and (ii)) the reset state observable.
Lemma 3.1. There exists well defined stationary distribution H. ((A- over histories likely to be seen at the limit of the interaction when मे is executed in P (proof in Appendix D) Remark 2. Assumption 2 is present in a vast majority of RL scenarios, where it is common practice to reset the environment and start anew from an initial state when the agent succeeds, fails, or after finite number of time steps (Brockman et al., 2016 Pardo al. 2018). A notable exception is in the domain of continual RL. However, the existence of a stationary distribution (Lem.B.j is often assumed in such scenarios (see, e.g., Huung 2020), which allows to relax the episodic assumption.
5

---

Published as a conference paper at ICLR 2024
Local losses. The objective function of the WAE-MDP incorporates local losses (Gelada et al. 2019)
that minimize the expected distance between the original and latent reward and transition functions:
LR 5,0,a~H4
[ a) oa)),, Lp
Wa (øPal- 1,,0 P(- ( (), a));
s,0,a~Hl
and both are optimized locally, i.e., under Hi where 8, O, a ~ He is a shorthand for (i) h ~ ai so that is the last observation of h, (ii)
T (*)), and (iii) a *(. (k)) [ Furthermore, p 8,0) a) is the distribution of transitioning to 8' ~ P(- 8, a). then embedding it to the latent space = (s"), and d is a metric on S. In practice, the ability of observing states during learning enables the optimization of those local losses without the need of explicitly storing histories. In- stead, we simply store the transitions of Ma encountered while executing म. We also introduce an observation loss which allows learning O:
Lo
E
E dTv
E
0(.1((6)0))
(4)
5,0,0~Hr P(-aa))
o'-O(1saa)
Belief losses. We set D as the Wasserstein distance between the true latent belief update and our be- lief encoder. In addition, we introduce the following reward and transition regularizers to reconcile the behaviors obtained in the fully observable latent model M and the partially observable one P:
L
E. h-
Wa (F" (h), x (h)),
L
E H. (-
E
R(6(s,0), o), - R(5,a)},
I P
E h, s (N)
E
Wa (P(- o). P(I1s,))))
(5)
L and LP aim at regularizing and minimize the gap between the rewards (resp. transition prob- abilities) that are expected when drawing states from the current belief compared to those actually observed. Again, the ability to observe the states during training enables optimizing those losses without explicitly requiring the states to execute the policy. The belief loss and the related two reg- ularizers be optimized on-policy, i.e., coupled with the optimization of the latent policy that is used to generate the episodes.
Value difference bounds. We provide guarantees concerning the agent behaviors in P. when the policies are conditioned on latent beliefs. To do so, we formalize the behaviors of the agent through value functions. For specific policy the value of a history is the expected return that would result from continuing to follow the policy from the latest point reached in that history: V. (h) Ex ["=00' rt by ** (h)]. Similarly, we write V for the values of the latent policy in P.
Suppose the agent uses a latent policy whose inputs are produced by ip, we claim that when the losses are minimized to zero, then (i) the latent model almost surely mimics the original environment, and (ii) our belief representation captures the value function. Precisely, let L = LR. + L + Lt +
y KK7 (Lp + IE L, + Lo), with R R and K7 R²/1 /1- then:
Theorem 3.2 (Model quality). For any latent policy /: B -- A(A). the values of P and P are bounded by the local and belief losses in average when, in both models, the actions are produced by which is conditioned on the latent belief induced by yp, i.e. a T( 11*1"*(0p (h)):
C
E 1-
V=())
V
(6)
Theorem 3.3 (Representation quality). Let 7 be an optimal policy of P, then for any > 0 and n & there a K > 0 so that for any histories hy h2 of length at most n that are measurable under P and P with up8 (hh)) = bi and p(12) b2, the representation induced by vp yields:
KWa b2), + + KL + L
1
V (h1) Vg. (h2))
+ Ha.
(7)
1-y
(11(
While Thm. 3.2 asserts that training a WAE-MDP as a latent space model of the environment results in similar behaviors (i.c., close expected returns) compared to the original environment when they are measured under the agent policy which justifies the usage of P as model of the environment Thm. 3.3 states that our learned update procedure yields a belief representation which well- suited to optimize the policy: execution traces leading close latent beliefs (via our learned updater <p) are guaranteed to yield close expected returns as well (proofs in Appendix F
Analogous we define yt the recursive application of yp along histories.
6

---

Published as a conference paper at ICLR 2024
B
C
t
t+1
Masked Stelo Flow
Belief
M
Aidden state
rub-taled encader
RNN
A2C Joes
A2C looas
A2C loss
A2C loss
Figure 2: WBU (leff) learns to encode the history into a sub-belief solely by optimizing the be- lief loss. The policy, being conditioned on the sub-belief, is learned via A2C and does not back- propagate through the sub-belief encoder. The R-A2C agent (right) uses BPTT: the RNN leverages gradients from future time-steps to improve its compression of the history for learning a policy and value function. In both plots, the colored arrows represent the gradient flows of the different losses.
4 LEARNING TO BELIEVE
In the following, we assume that we have access to the latent model learned by the WAF-MDP.
Architecture. Our latent belief encoder aims at generalizing to any POMDP. Therefore, we do nof make any assumption about the underlying belief distribution. To accommodate complex belief distributions, we use a Masked RegreRegees Flows (MAF) (Papamakarios et al. 2017). a type of normalizing flow built on the auto-regressive property. Precisely, to fit with the WAE-MDP framework and leverage the guarantees presented in Sect. 3.2 we use the MAF of Delgrange et al. (203) that learns relaxed multivariate latent distributions.
The sub-belief 3r is the vector that embeds the parameters of the belief distribution, which is con- verted into belief via the MAF M(3,) = be. We use a enccoee encoder shth to recursively update
3e via (B,, de, 0(+1)= = 3.+1, so that yp (be, at, 0t+1) M (Be, G2, o++1) RNNs are trained via back-propagation through time (BPTT), which is challenging (Pascanu et al. 2013). In contrast, albeit sub-beliefs are updated recursively in the same spirit as RNN hidden states, we do not need to use BPTT and use a sieplefeed-fooe feed-forwate network for qdth, as illustrated Fig. 2 In R-A2C, RNN hidden states serve as compact representations of histories for the policy. Since values of time-steps closer to the end of an episode are easier to learn, the gradients of future time-steps tend to more accurate; thus BPTT helps learning. This is in stark contrast to learning the belief update rule: the beliefs of early time-steps are casier to infer, so BPTT is unnecessary and might even be armful. Additionally, viewing the belief update function as the transition function of MM. disabling BPTT aligns with the literature on model-based RL for learning Markovian transition functions (Gelada al. 2019 François-Laver et al. 2019). See Appendix
G for a more detailed discussion on BPTT.
Training. We aim to train yh and M to approximate the update rule by minimizing the Wasserstein between the belief update rule F of the latent POMDP, and the belief encoder yp (E.. to leverage the theoretical learning guarantees of Thm. and 3.3 However, Wasserstein optimization is known to be challenging, often requiring the use of additional networks, Lipschitz constraints, and min- max optimization procedure (Arjovsky et al. 2017). similar to how WAE-MDPs are trained. Also, sampling from both distributions is necessary for optimizing Wasserstein and, while sampling from our belief approximation is straightforward, sampling from the update rule (Eq. 3) is non-trivial.
As an alternative to the Wasserstein optimization, we minimize the KL divergence between the two distributions. While this DKL. proxy is easier to optimize and only requires sampling from one of the two distributions (in our case, the belief encoder), it bounds the Wasserstein distance by the Pinsker's s inequality (Borwein& & Lewis 2005) in the pero-temperature limit of the WAE-MDP (cf. Appendix for a discussion).
On-policy KL divergence. Using DkL, as proxy for W. allows to narrow the gap between şo and We train y on-policy, with the same samples used for म, which aids learning despite gradients not allowed to flow between the networks, At any time-step & > O, given the current belief b,
7

---

Published as a conference paper at ICLR 2024
Repeet Previous
1.0
Cartpole
Noisy Stateless Cartoole
Comulatius roward
0.0
Comuletive reward
66
0.5
0.G
Cumulatve
0.4
0.4
0.2
0.
00.
40
100
0
40
60
100
0.0
40
60
100
(x10000)
(X10000)
(x10000)
20
Space Invaders
15.0
Notey Space invaders
10.0
Bellet loss
Comuzativn reppard
Cunulative relward
12.5
15-
Neay Saetalca Cartzoli
20.0-
Bellct toss
5.0
Space fraccess
Nolyy Space twaders
7.5
2.5
2.5
-2.5
50
200
0.0
50
150
200
100
150
200
Steps (110000)
Steps (110000)
Steps (x10000)
WBU fours]
R-A2C
DVAL
Figure 3: Evolution of the (i) undiscounted cumulative return for WBU, R-A2C and DVRL, and (ii) estimated belief loss during learning for WBU (mean and standard error). We report instances of each algorithm. Appendix Appendix[detaill, details the hyperparameter search performed.
the action o, played by the agent, and the next perceived observation or+1, the belief proxy loss is:
DKL (p (b,, at, 0c+1) 7 (be, at, at, oc+1)) log E 8mbb 3',P(a),.e))
T (oct1 8)))
log sp (5e41 bi, az, 0(41) log E. P( 5, 5, - log (oc+1 Bet1) 8mbb
E
(8)
Rp (14))
Eq. 8 consists of 4 terms: normalization factor, negative entropy of yp belief update conformity with the latent MDP's state transition function, and filtration of latent states unrelated to Or+1- Remark 3 (Latent observations). The WAE-MDP learns from the augmented POMDP pt (Sect. (3.1 which is equivalent to P and possesses deterministic observation function. The WAE-MDP learns such a deterministic mapping through O,,. However, when deterministic, the observation terms of Eq. 3 and 8 are Dirac, which prevents learning: deterministically filtering out states from the next belief that do not share the next observation is sample inefficient and may further require construct- ing the full belief distribution, which is usually intractable. In practice, to alleviate this, model the observation function as a normal distribution: O- 8) N(OO,, (8), o²) where the variance can, recover a Dirac.
e.g., be learned via Lo. Notice that we can enforce annealing the variance
Policy learning is enabled by inputting the sub-belief into the policy, while the optimization of the belief encoder parameters by the RI. agent not allowed. Our method is applicable to any on-policy algorithm, and we employ A2C in our experiments. We provide the final algorithm in Appendix H
5 EXPERIMENTS
To evaluate our approach, we identify three types of POMDPs: those requiring long-term memory, those where features of the state space hidden (and may be inferred from short-term memory), and those with noisy observations. Notably, we stress that long-term memory is crucial in POMDPs, whereas short-term memory could be mitigated by stacking frames (e.g., Mnih al. 2015). We compare our agent to R-A2C and DVRL (Fig. 3) trained in environments from POPGYM MMoad al. 2023) and our own partially observable version of MINATAR (Young & Tian 2019).
Memorization. The REPEATPREVIOUS environment involves shuffling two decks of cards at the start of each episode and presenting the agent with card at each time step. The goal is to identify the suit of the card seen time steps earlier. Our algorithm stands out as the sole method demonstrating mid- to long-term memorization capabilities. Unlike other methods, notably DVRL which also attempts to learn a belief distribution, WBU provably acquires a suitable representation of the history by learning to maintain sufficient statistic, thereby explaining its ability to retain past information.
8

---

Published as a conference paper at ICLR 2024
Hidden features. We employ cart pole scenario (STATELESSCARRTPLLE where velocity compo- nents of the system are hidden. R-A2C excels rapidly here, capitalizing on short-term memory to infer velocities from the preceding observation, while DVRL is overtaken. Still, WBU eventually reaches R-A2C final performance. We also explore the SPACEINVADERS environment, where the agent takes command of cannon with the objective of shooting at groups of moving aliens. In the observation, we intentionally concealed the direction of alien movement and confounded friendly and enemy fires. In this more challenging setting, WBU excels by earning the highest rewards.
Noise. We explore two types of noise. First, we introduce Gaussian noise to the observations of STATELESSCAKTPOLE Second, for SPACEINVADERS, binary noise injected via a radar-like mask obscuring the position of each alien with high probability. Hence, the agent must infer their positions based on previous observations. By leveraging ability to maintain belief over the noiseless (latent) state space, WBU demonstrates its resilience to noise and swiftly provides superior solutions, whereas R-A2C eventually achieves comparable performance but with more variance.
Belief representation. Ideally, close policy inputs should lead to close values, which would case its optimization. Thm. 3.3 provides such represen- tation guarantee and ensures that the representation induced by & captures the value function. To sup- port this, we performed a t-SNE (van der Maaten & Hinton 2008) on our belief representation the late stage of training, which projects latent beliefs on a 2D space (Fig. 4. Interestingly, latent beliefs clus- tered together have indeed close values, in line with Thm. 3.3 We reported the belief loss throughout the training phases (Fig. 3) Importantly, unlike other baselines, our approach distinctly separates the op- timization of up and न. Consequently, the policy opti- mization does not influence the representation which is solely learned via ip The decrease in this loss thus relates to improved representation quality for RL.
Value
Figure 4: Two-dimensional t-SNE of our be- lief representation for SPACHINVADERS
6 CONCLUSION
WBU provides a principled approach that approximates directly the belief update for POMDPs, in contrast to SOTA methods that uses the RL objective and regularization to attempt to turn the history into a sufficient statistic. By learning the belief and its update rule, we provide strong guarantees on the quality of the belief, its ability to condition the optimal value function, and ultimately, the effectiveness of our algorithm. Our theoretical analysis and experimental results demonstrate the potential of our approach. Overall, WBU provides a promising new direction for RL in POMDPs, with potential applications in a wide range of settings where decision-making is complicated by uncertainty and partial observability, when guarantees the agent behaviors are required.
Future work. The theory we developed is not limited to the algorithm proposed in our paper. It opens diverse avenues for future work, e.g., on formally verifiable policies for POMDPs, by lever- aging the guarantees presented in our framework. We also leave the study and the adaptation of our framework under relaxed assumptions (e.g., in settings akin to the work Lambrechts et al. 2022) to future work. In addition, Thm 3.2 enables policy optimization through planning in the learned model, as demonstrated successful model-based RI. methods (e.g.. Hattner et al. 2021 Scaling to nigh-dimensional observations (e.g.. images) may potentially be computation-intenset observation filtering. For this further challenge, we suggest either to modify the WAE-MDP frame- work by using a stochastic decoder (see Tolstikhin et al. 2018 e.g., via PixelCNN, van den Oord et al. 2016), or learning a diwer-dimensional latent observation space synced with the policy with a normalized or (relaxed) discrete prior (e.g., via a WAE-GAN Tolstikhin et al. 2018). Finally, incor- porating bisimulation metrics (Desharnais et al.) 2004 Ferns et al. (011)) will strengthen guarantees for belief learning, even though bisimulation is challenging in POMDPs (Castro al. 2009
due to

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This research was supported by funding from the Flemish Govermment under the "Onderzoekspro- gramma Artificiële Intelligentie (AI) Vlaanderen" program and was supported by the DESCARTES iBOF project. R. Avalos is supported by the Research Foundation Flanders (FWO), under grant number 11F572IN. G.A. Perez is also supported by the Belgian FWO "SAILor" project (G030020N). We thank Mathicu Reymond, Denis Steckelmacher, and Mustafa Mert Çelikok for their valuable feedback.
REPRODUCIBILIT STATEMENT
We referenced in the main text the parts of the Appendix presenting the proofs of our Lemma (Ap- pendix D) and Theorems (Appendix F.. We also provide the pseudo-code of our algorithm (Ap- pendix[H] HD as well as extra details required to compute our losses (Appendix [ E and The code
available at atttps : / /tt/thuuubb/.// com raphae ava l os owbbb Additionally, we provide the details of our hyperparameter search (Appendix
REFERENCES
Martín Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial net- works. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Con- ference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, vol- 70 of Proceedings of Machine Learning Research, pp. 214-223. PMLR, 2017. URL
http: http:://proeeedings mlr www.pppwsssssee.e / html
Raphael Avalos, Mathicu Reymond, Ann Nowé, and Diederik M. Roijers. Local Advantage Net- works for Cooperative Multi-Agent Reinforcement Learning. In AAMAS '22: Proceedings of the 2 lst International Conference on Autonomous Agents and MultiAgent Systems (Extended Ab- stract), 2022.
J.M. Borwein and A.S. Lewis. Convex Analysis and Nonlinear Optimization: Theory and Examples. CMS Books in Mathematics. Springer New York, 2005. ISBN 9780387295701. URL https: //books. , gwogle.ooooooooooooooooooo btwp/ookksi...xx...e......
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. CoRR, abs/1606.00 5400 2016. URL httl:/tttaxxxi..oo/ /rgg
abs/1606 01540
Peter E. Bulychev, Franck Cassez, Alexandre David, Kim Guldstrand Larsen, Dean-François Raskin, and Pierre-Alain Reynier. Controllers with minimal observation power (application to timed sys- tems). In Supratik Chakraborty and Madhavan Mukund (eds.), Automated Technology for Veri- fication and Analysis 10th International Symposium, ATVA 2012, Thiruvananthpurrmm India, October 3-6, 2012. Proceedings, volume 7561 of Lecture Notes in Computer Science, pp. 223- 237. Springer, 2012. doi: 10. 1001/9077 0-777833-62---3 04442-3336-6,,00 URL https:://oi.croetttt
1007 978 842 3338 19
Pablo Samuel Castro, Prakash Panangaden, and Doina Precup. Equivalence relations in fully and partially observable markov decision processes. Craig Boutilier (ed.), IJCAI 2009, Proceedings of the 2t International Joint Conference Artificial Intelligence, Pasadena, California, USA, July 11-17, 2009, pp. 1653-1658, 2009. URL http::/ijcaa..rrg////// wtt.://cece...e.eeeeeeeeeeeeee/e Papers/277 Papers/276.por
Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. Mico: Im- proved representations via sampling-based state similarity for markov decision processes. In Mare'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wort- man Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Confer- Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 301 30126,, 2021. URL httpss////rooeeeinnggg / 11t.a//eee..0e.....2e22ee1111110
www.eeiiiiiinnn.ni......... /ppp////// htt7://a//2000000000022211223113
10

---

Published as a conference paper at ICLR 2024
Krishnendu Chatterjee, Luca de Alfaro, Rupak Majumdar, and Vishwanath Raman. Algorithms for game metrics (full version). Log. Methods Comput. Sci,, 6(3), 2010. URL http::
org/ abs/ harg/abs/oe 0809.. 4326
Krishnendu Chatterjee, Martin Chmelik, Raghav Gupta, and Ayush Kanodia. Optimal cost almost- sure reachability pomdps. Artif. Intell., 234:26-48, 2016. doi: 10. 2..1066).artimt2216.01....
URL httph: //t:: wwooi.org/1i wtw.:...1....a....
2016 6. 01 007
Xiaoyu Chen, Yao Mark Mu, Ping Luo, Shengbo Li, and Jianyu Chen, Flow-based recurrent belief state learning for pomdps. In International Conference on Machine Learning, pp. 3444-3468. PMLR, 2022.
Luca de Alfaro, Thomas A. Henzinger, and Rupak Majumdar. Discounting the future in systems theory. In Jos C. M. Baeten, Jan Karel Lenstra, Joachim Parrow, and Gerhard J. Woeginger (eds.), Automata, Languages and Programming, 30th International Colloquium, ICALP 2003, Eindhoven, The Netherlands, June 30 July 4, 2003. Proceedings, volume 2719 of Lecture Notes in Computer Science, pp. 1022 1037. Springer, 2003. doi: httpss:/ddot htts://doi.org////// 1111111.110010000 00110000033-00-0000 5-00540040000 45061 1021110077
URL
Florent Delgrange, Ann Nowé, and Guillermo A. Pérez. Distillation of rl policies with formal guar- antees via variational abstraction of markov decision processes. Proceedings of the AAAI Con- ference on Artificial Intelligence. 36(6):497765555, Jun. 2022. doi: 10.1609/amai, v36.6.2066666
URL https:
org/ wtog/index..m
Florent Delgrange, Ann Nowe, and Guillermo Perez. Wasserstein auto-encoded MDPs: Formal ver- ification of efficiently distilled RL policies with many- guarantees. In International Confer- ence on Learning Representations, 2023. URL https::t////..... /woperpeeeeveeeeeeeeeeeeeeee net/ forum?
httpps://openrevieeee
JLLTT Ech12Y
Josée Desharnais, Vincet Gupta, Radha Jagadecsan, and Prakash Panangaden. Metrics for labelled
markov processes. Theor. Comput. Sci., 318(3)3333333333 354, 2004. doi: (0.10166..es.2000.0..011
URL http:p/s//do...oggg1.000000
wwt.:ttssscoo.oooooooooooooo 09.013
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. S/AM J. Comput., 1066)1661111114,, 2011. doi: 10.1137/10080484X. URL
otttp::/ddo..org wtt.:////00011111100000000000000
Vincent François-Lavet, Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined reinforce- ment learning via abstract representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3582-3589, 2019.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofrr Nachum, and Marc G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2170-2179. PMLR, 2019. URL. /RRphhttp://proceeeeee mlr www.pess/v9..
ada 9a. html
Robert Givan, Thomas L. Dean, and Matthew Greig, Equivalence notions and model minimization in markov decision processes. Artif. Intell., 143--22) 2003. doi: 00016/0000030022020 00376-4. httts:t:: org.iiiorr/ooo 10. 0010////00000000 3702 (C2) 00376-4
Karol Gregor, George Papamakarios, Frederic Besse, Lars Buesing, and Théophane Weber. Tem- poral Difference Variational Auto-Encoder. 7th International Conference on Learning Represen- tations, ICLR 2019, 2018. doi: (tt.5//aaxxxi.11003331000 URL https:://arxiv.. abs/ 1806. 03107v3
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aäron van den Oord. Shaping Belief States with Generative Environment Models for RL. Advances in Neural Information Processing Systems, 32, 6 2019. ISSN 10495258, doi: 0..48850/axiiy.1966699 htl.:ts:: /wt//axiiia org abs 906. wttpp///.1166600...2.0 09237v2

---

Published as a conference paper at ICLR 2024
Danijar Hafner, Timothy Lillicrap Deepmind, Jimmy Ba, Mohammad Norouzi, and Google Brain. Dream to Control: Learning Behaviors by Latent Imagination. 12 2019. doi: 10.48550/arxiv. 1912.01603. URL
abtp:///aaaa/aaa 01603v3
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021. URL
https: /pwp://veeeeeeeeeeeeeemmmeeeeeee net/ forum? diw.ummmddd00aayy200000
Eric A. Hansen. Solving pomdps by searching in policy space. In Gregory F. Cooper and Serafín Moral (eds.), UAI '98: Proceedings of the Fourteenth Conference on Un- certainty in Artificial Intelligence, University of Wisconsin Business School, Madison, Wisconsin, USA, July 24-26, 1998, pp. 211-219. Morgan Kaufmann, 1998. URL.
jap ?mmnu cocennullsmmmnn 1&smnu
2tarticle_id id- htt.riicce_.d-220ppprrceeedigeed
Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable MDPs, In AAAI Fall Symposium Technical Report, volume FS-15-06, pp. 29-37. Al Access Foundation, 2015. ISBN 9781577357520.
Milos Hauskrecht. Value-function approximations for partially observable markov decision pro- cesses, J. Artif. Intell. Res., 13:33-94, 2000. doi: 10.161 10.1613/jair.675 URL httlttt////o.......
10.1613/ 10r.163//air... 668
Bojun Huang. Steady state analysis of episodic reinforcement learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020. URL httpss:/prrocedinggs iew.rrrreeednnss...ei.ssssss neurips. cc/pagger2222/hash. 5901242221135f5551000aaaaaa2-2
.ttll
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for POMDPs. In 35th International Conference on Machine Learning, ICML 2018, volume 5, 2018.
Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134 1998.
Ivan Kobyzev, Simon J.D. Prince, and Marcus A. Brubaker. Normalizing flows: An introduction and review of current methods. IEEE Transactions on Pattern Analysis and Machine Intelligence, 43(11):3964 3979, nov 2021. doi: 10. http:/pami.222..9992222 URL CRLhttts. https: http::::///oooo///// wtw.:aaaaa.......mmmm
10. 10.110992Ftpami 2020 2020.299293.3
Gaspard Lambrechts, Adrien Bolland, and Damien Ernst. Recurrent networks, hidden states and beliefs in partially observable environments. Transactions on Machine Learning Research, 2022.
ISSN 2835-8856. httptt/ppss///eeeeeeeeeee net wttttfeumm/eeeeee (dforumm?dd=dktft33e3......
Gaspard Lambrechts, Adrien Bolland, and Damien Ernst. Informed POMDP: information in model-based RL. In ICML Workshop on New Frontiers in Learning, Control, and Dynamicul Systems, 2023. URL
Leveraging additional
http://ett//ouummooom/oo form//orrmmmmm F=mm4 sF sFookkkk
net.
Kim Guldstrand Larsen and Arne Skou. Bisimulation through probabilistic testing. In Con- ference Record of the Sixteenth Annual ACM Symposium Principles of Programming Lan- guages. Austin, Texas, USA, January 11-13, 1989, pp.
ACM Press, 1989. doi:
10.1145/75277.75300
hnttps:/:///o org
75307
Mikko Lauri, David Hsu, and Joni Pajarinen. Partially observable markov decision processes in robotics: A survey. IEEE Transactions Robotics, 39(1):21-40, 2023. doi: 10.1 N0.1109/TR0.002. 3200138.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Advances Neural Information Pro- cessing Systems, 33:741-752, 2020.
12

---

Published as a conference paper at ICLR 2024
Jonathan Lee, Alekh Agarwal, Christoph Dann, and Tong Zhang. Learning in pomdps is sample- efficient with hindsight observability. In International Conference on Machine Learning. pp- 18733-18773. PMLR, 2023.
Xiao Ma, Peter Karkus, David Hsu, and Wee Sun Lee. Particle filter recurrent neural networks. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 5101-5108, 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonogiou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529-5333 533, 2 2015. ISSN 14764687. doi: 10. 1038/nature14220
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous Methods for Deep Reinforcement Learning. 33rd International Conference on Machine Learning, ICMI. 2016, 4:2850-2869, 2016. httttt//////xiii www.oorgg/aaaaooooooooooo.oooaaa 602. C1783
Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POPGym: Benchmarking partially observable reinforcement learning. In The Eleventh International Confer- ence on Learning Representations, 2023. URL ltttps://openzee////
net/ wwrummiiin
chDrutuUTsCK]
Yu Nishiyama, Abdeslam Boularias, Arthur Gretton, and Kenji Fukumizu. Hilbert Space Em- beddings of POMDPs. In Proceedings of the Twenty-Fighth Conference on Uncertainty in Artificial Intelligence (UA12012), 10 2012. doi: 10.485500amxxx.1110.4463 URL. https:
Frans A. Olichoek, Matthijs T.J. Spaan, and Nikos Vlassis. Optimal and approximate Q-value functions for decentralized POMDPs. Journal of Artificial Intelligence Research, 32:289-353, 2008. ISSN 10769757. doi: 10. 10.1613/jair:244.. URL http: dei. oro/.0..1111111 jair 2447.
George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density timation. In I. Guyon, U. Von Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran As- sociates, Inc., 2017. URL
neurips pwtppppppe//1////lee....
6clca886822 020222228882220077000085555--000
pdf
Fabio Pardo, Arash Tavakoli, Vitaly Levdik, and Petar Kormushev. Time limits in reinforcement learning. In Jemnifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmüssan, Stockholm, Sweden, July 10- 15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 4042-4051. PMLR, 2018. httl:tttp/pooceeeeee
www.ppewssssepppppppppppp patdo Ba html
Razvan Pascanu, Tomás Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In Proceedings of the 30th International Conference on Machine Learning, ICML 2013, Atlanta, GA, USA. 16-21 June 2013, volume 28 of JMLR Workshop and Conference Proceed- ings, pp. 1310-1318. JMLR.org, 2013. URL http://prooeedinga mir mnlr.peess/28/ pascanu1 3. html
Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. Wi- ley Series in Probability and Statistics. Wiley, 1994. ISBN 978-0-47161977-2. doi: 10.1002/ 9780470316887. URL htttps:://c...rgg/// 11t.1....100000001110000
6887
Sebastian Raschka, Joshua Patterson, and Corey Nolet. Machine learning in python: Main develop- ments and technology trends in data science, machine learning, and artificial intelligence. arXiv preprint wwXii:2002.04803, 2020.
Richard D. Smallwood and Edward J. Sondik. The optimal control of partially observable markov processes over a finite horizon. Oper. Res., 2165):1071-1888 1973. doi: 10. 1287//ppe221.551107. URL mnttps://dos www.doi..oog/oooo orgg
1071
13

---

Published as a conference paper at ICLR 2024
Edward J. Sondik. The optimal control of partially observable markov processes the infinite
horizon: Discounted costs. Oper. Res., 26(2):282 304, 1978. doi: 10.12887/ome:2622.222.. URL
httpps:/doi. https://coo..ccg//...22211/
26.
R.S. Sutton and A.G. Barto. Reinforcement Learning: An Introduction. IEEE Transactions on Neural Networks, 1998. ISSN 1045-9227. doi: 10. 100/tnn.1999.. 7929992
Ilya O. Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schölkopf. Wasserstein auto- encoders. 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018.
URLltttps / http///ppppeeeeeeeeee/ee/eeeeeee net wwt.:nett//tiiiiooom /tt////rrrmmmmeeeee/mm iidmm?dddkkiiii 0b
Aäron van den Oord, Nal Kalchbrenner, Lasse Espeholt, Koray Kavukcuoglu, Oriol Vinyals, and Alex Graves. Conditional image generation with pixelenn decoders. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 29: Annual Conference on Neu- ral Information Processing Systems 2016. December 5-10, 2016. Barcelona, Spain, pp. 4790-4798, 2016. URL /ttpststss/ htwp://ceeeiiag................e neurips hccppaper/2cl//haeeee b1301 41 fef Fow -bst
html
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Ma- chine Learning Research, 9(86):2579-260) 2008. URL http: jmlr wttppppprppapps//////
wandermaatennSa. html
Elise van der Pol, Thomas Kipf, Frans A. Oliehoek, and Max Welling. Plannable approximations to MDP homomorphisms: Equivariance under actions. In Amal El Fallah Seghrouchni, Gita Sukthankar, Bo An, and Neil Yorke-Smith (eds.). Proceedings of the 19th International Confer- ence on Autonomous Agents and Multiagent Systems, AAMAS '20, Auckland, New Zealand, May 9-13, 2020, pp. 1431-1439. International Foundation for Autonomous Agents and Multiagent URL
Systems, 2020. 3398926
5555/339 doi: 10. (0.5555/3398761.33 033999997661331966..
acm. wttp://a..ccoooooo/ooo/o/ooooo//
Cédric Villani. Optimal Transport: Old and New. Springer Berlin Hcidelberg, Berlin, Heidelberg, 2009. ISBN 978-3-540-71050-9,
doi: 10.1007/978-3-540-70000 6. URL /ttop: org/10. 1007 1t11:/1000077772--------------00
Yuhui Wang and Xiaoyang Tan. Deep recurrent belief propagation network for pomdps. In Pro- ceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10236-10244, 2021.
Yaxiong Wu, Craig Macdonald, and Iadh Ounis. Partially observable reinforcement learning for dialog-based interactive recommendation. In Proceedings of the 15th ACM Conference on Rec- ommender Systems, pp. 241 241, 2021.
Kenny Young and Tian Tian. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv preprint arXiv: 2019.
Hongyu Zang, Xin Li, and Mingzhong Wang. Simsr: Simple distance-based state representations
for deep reinforcement learning. In Thirty-Sixth AAAI Conference on Artificial Intelligence, AAAI 2022, Thirty-Four Fourth Conference on Innovative Applications of Artificial Intelligence, IAAI 2022, The Twelveth Symposium on Educational Advances in Artificial Intelligence, EAAI 2022 Virtual Event, February March 1, 2022, pp. 8997-9005. AAAI Press, 2022. doi: 10. AA09AA.
V3618.20883. URL. https: / io www.o/.gg.oooooo 609 /t////aaa.... v3618.8..0 2e0.68.0088.0 20883
Tianhao Zhang. Yueheng Li, Chen Wang, Guangming Xie, and Zongqing Lu. FOP: Factorizing Optimal Joint Policy of Maximumm-Entropy Entropy Multi-Agent Reinforcement Learning. In Proceedings of the 38th International Conference on Machine Learning, volume 139, 2021.
K.J Aström. Optimal control of markov processes with incomplete state information. Journal of Mathematical Analysis and Applications, (11)::174205, 1965, ISSN 0022-247X. doi: http:://db.ogg//l..0000002222211 http:///..aaooorrrr.ee//////2222
URL https: /wwww sciencedirect.
14

---

Published as a conference paper at ICLR 2024
APPENDIX
A ADDITIONAL RELATED WORK
In the same spirit as in our work, the observability assumption of Lee et al. (2023) aligns perfectly with ours, but is leveraged in a different way, to provide results oriented towards sample complexity in the tabular case, as well as regret bounds when particular types of function approximators are used. Several complexity results and regrets bounds were developed for sub-classes of POMDPS in finite horizon, such as when the observability is well- (weakly revealed POMDPs, Liu et al. 2022), or when the sequential decision making problem admits low rank structure (Liu et al. 2022 Huang et al. (222). On the contrary, our results hold for general POMDPs; precisely, our constraints rather on the training procedure and not on the nature of the POMDP. Finally, we note that RQL- AIS (SeyedSalehi et al. 2023) presents analysis of recurrent Q-learning for the tabular setting which also includes value difference bounds where their latent state is the (discrete) hidden state of an RNN. While theoretically interesting, their bounds are generally intractable as they rely on global metrics requiring them to be minimized over the whole history and action space. This limitation is made clear in their algorithm as they use the local version of those metrics as regularization losses. In our case, because we focus on-policy metries (local losses) our theoretical guarantees hold in practice.
Bisimulation. One of our result involves abstraction quality guarantees (Theorem 3.2) formalized as the closeness of the value function of the original and latent models. In spirit, this result can be related to bisimulation. Bisimulation is a behavioral equivalence relation between states and mod- els whose properties include trajectory and value equivalence (Larsen & Skou [999; Givan et al., 2003). However, due to their "all or nothing" nature, bisimulation relations are often too restrictive; intuitively, states that share close rewards and dynamics but exhibit slight numerical differences are treated as being completely different. Bisimulation pseudometrics (de Alfaro et al.) 2003 Deshar- nais et al. 2004) generalize bisimulation by assigning a distance between states and quantifying their bisimilarity instead of yielding binary signals indicating whether they are bisimilar or not. Specifi-
cally, realev (discounted) quantitative properties (de Alfaro et al. 2003] Chatterjee et al. 2010)
as well as expected returns (Ferns al., 2011) are guaranteed to be Lipschitz continuous w.r.t. the
related bisimulation distance (compare with the almost- Lipshssttssssssssssssss
our Theorem 3.3).
Works representation learning for RL (c.g.. Gelada et al. 2019 van der Pol ct al. 2020 Castro
et al. 2021 Zhang et al. 2021: Zang et al., [022] Delgrange et al. 2023) use bisimulation metrics to optimize the representation quality offered by latent space, and aim learning representations which capture the bisimulation distance, so that two states close in the representation provide close and relevant information for learning the policy. Precisely, such works seek the Lipschitz continuity of the optimal value function with regards to the distance between pairs of points in the latent space.
Note that while our work builds upon the WAE-MDP framework (Delgrange et al. 2023). which learns a latent model based on bisimulation metrics, our work does not learn any bisimulation- like metric in the underlying latent space of the POMDP. Instead, we prioritize what truly matters in the context of representation for RL: the Lipschitz
of the optimal value function. In particular, the attentive reader may notice that, in contrast to usual representation bounds in MDPs (e.g.. Gelada et al.] 2019 Delgrange et al.] [0022), our bound involves the almost Lipschitzness the value function, as well as various additional losses linked to the belief update and observation functions, which underlines the challenge of providing such guarantees. Indeed, bisimulation in the context of POMDPs represents a major challenge (e.g.. Castro et al.] (009):: intuitively, in the context of POMDPs, the bisimulation equivalence is not defined on the state space but rather the belief space induced by the different components of the POMDP. Notably, even if the POMDP is finite, i.e., state, action, and observation spaces are finite, the belief space can be represented as the usual (uncountable space) simplex, which prevents the application of bisimulation algorithms developed in the context of countable MDPs.
B THE BELIEF UPDATE RULE
Let P = (M, S2 o) be a latent POMDP with underlying MDP M = 51,7). At step t > 0, assume that the current latent belief is by B. Then, when a A is executed and 0t+1 E Ω
15

---

Published as a conference paper at ICLR 2024
is observed, by is updated according to 7(b,, az o(+1) as follows, that for any (believed) state St+1 8,
(5+1(5 Bs ~h, P(5++ 5e, ac) O Se+1)
(9)
Ber E g
( 88)
(,a))
B P(a s..
Bacc P(s i.
S (ot+:
b++1
R
Figure 5: The belief update rule: (left) transformation of the current belief b, with the transition probability function P. evaluated the current action or, into the state probability density; (right) filtering out the next states that could not have produced the next observation ot-1-
The belief update rule divided in two steps (cf. Fig. First, the current belief distribution bt is used to marginalize the latent transition function P over the believed latent states, to further infer the distribution over the possible next states. This first part corresponds to computing the probability of going to the next believed state 8141 from the possible states 8, witnessed by the current belief b, Second, the next observation 0++1 is used to filter the resulting density based on the observation function: believed latent states that do not share the observation o++1 perceived filtered out according to (oc++ 1)) It is worth noting that the latent model is learned from Pi, whose observation function is deterministic. Without modelling the latent observation function O normal distribution, the second part of the belief update would need to eliminate all next states with different observations - which is not gradient descent friendly. The third operation (not present in Fig. 5) normalizes the output of the observation filtering to obtain probability density.
C DIRAC MEASURES
In this work, we consider the Dirac delta function ê as a measure. Specifically, this means that for any complete, separable space and point E X' the Dirac measure with impulse is 8, A(X) and satisfies (A) = if a E A and (A) 0 otherwise, for any A E E(X). Interesting properties of the Dirac measure include 6, limm- -o Naa, o²), where ((a, o²) mean a and variance o², and Sx (z) dr f(a) for any compactly supported function f.
) is the normal distribution with
D PROOF OF LEMMA STATIONARITY OVER HISTORIES
Let us formally restate the Lemma:
Lemma D.1. Let be an episodic POMDP with action space A and observation space S2 and т be a latent policy conditioned on the beliefs of latent POMDP P that shares the action and observation spaces of while being lifted P via the belief encoder, i.e., a T(- xpt (h)) for any E (A Q)* Then, there a well defined stationary distribution H₂ A ((A 2) over histories drawn at the limit from the interaction of the RL agent with P, when it operates under the latent policy n.
Proof sketch. Build history unfolding as the MDP whose state space consists of all histories and keeps track the current history of P at any time of the interaction. The resulting MDP remains episodic since it is equivalent to P: the former mimics the behaviors of the latter under म. All episodic processes are ergodic (e.g., (0uang2020) which guarantees the existence such a distri- bution.
16

---

Published as a conference paper at ICLR 2024
We dedicate this Section to formally detailing and proving every claim of this proof sketch. Before going further, we formally define the notion of episodic process, and we further introduce the notions of memory- based policies, Markov Chains, and limiting distributions in Markov Chains
D.1 PRELIMINARIES
Recall by Assumption 2 that the environment P is an episodic process. formally recall the notion of episodic process:
Definition D.2 (Episodic RL process). The RL procedure is episodic iff the environment embeds special reset state Sreset S so that
(i) under any policy и. the environment is almost surely eventually reset:
[P Some: Og 3t >> 8t Sreset })
(i) when reset, the environment transitions to the initial state: P(ar Sreset, a) > 0 and P(S)
} Sreset, a) for all A, and the distribution is the same for all ac- tions, for all a, a' E A, P(- Sreset: a) = P(- Sreset: d');
(iti) the reset state is observable: there is an observation o* Qt so that O(o" | |s'a) = when
s' + Sreset, and O-. Sreset, a) dor for E A.
An episode history (do:r-1; Ou.)) where O(1 (81,00) and or o".
Encoding memory through Mealy machines. Policies are building blocks to define the probabil- ity space any MDP. To deal with policies whose decisions are based on unrolling histories, we formally define the notion of memory in policies.
Definition D.3 (Policy as Mealy Machine). Given an MDP M (S, R. $1, ?) any policy for M can be encoded as a stochastic Mealy machine π = (Q, Ta, Tp. 91) as follows: Q is a set of
memory states; Tr S Q - ((A) is the next action function; TH S x Q x A the memory update function; and qı is the initial memory state.
A(Q) is
Example (Stationary policy). A stationary policy π can encoded as any Mealy machine π with memory space Q where [Q] 1. Example (Latent policy). Let P (M, O) with underlying MDP M (S, A,, and the latent space model P with initial state 5 be the POMDPs of Lemma D.1 Then, any latent (stationary) policy /: B ((A) conditioned on the belief space B of P can be executed in the belief MDP MR of via the Mealy machine n' (B, Ta. uu S, keeping track in its memory of the current latent belief E inferred by our belief encoder yp This enables the agent to take its decisions solely based on the latter: a. b, b) .. b). When the belief MDP transitions to the next belief b', the memory is then updated according to the observation dynamics:
Bash Pe-e(() O01e" .0) " (B') (7(h,a,))))) if * 8
(b 1 (,,,,,b))
Bass By P(-lea)) Eo Bo--o(-o (7(b),.0))(1)
T- b, b, Saica) 8 that
otherwise (to fulfil the episodic constraint), Fu is simply obtained by applying the usual conditional probabil- where Pr(V', 1 (,,,,,, b)
Note
ity rule: Na b, (,,,b,,,1)))
Pr(E 11b.B// E.-h P(-ss.) (o- 6 (b),,.//)))) (6) o-(b,a,o))()) and Pr(U/ b, b.
Ps( b,a)
since the next original belief state is independent of the current latent belief state.
Definition D.4 (Markov Chain). A Markov Chain (MC) is MDP whose action space A consists of a singleton, i.e,, |A| = Any MDP M = (S, R. 81,7) and memory-based policy π = (Q, Tas nu. 91) induces Markov Chain
M (S Q Pr, R. (81,91) (1,,)))))
where:
2 for convenience, we use the next state as input of the memory update function in addition to the current state of the MDP. This is non-standard but we contend that this change is harmless. Indeed, one may construct an MDP with dummy state (s, a, 8²) which is entered each time the system transitions from to s' when is chosen, before going to s'. Then, it suffices to change the discount to / to compensate for the addition of this transition.
17

---

Published as a conference paper at ICLR 2024
the state space consists of the product of the original state space and the memory of n the transition function embeds the next action and the policy update functions from the policy, i.e.,
P ((8,, ((88,g0 (((,,''' (8,9))
E Th Tp 8, (,,,a,8) P5 8, s,a), and
the rewards are averaged over the possible actions produced by the next action function, i.e,, R.,((s.))) Eg-na(laa R(s,a).
The probability measure I-M is actually the unique probability measure defined over the measurable
infinite trajectories the MC M (Baier & Katoen 2008). We omit the memory dimension when is stationary (i.e., memoryless).
I Limiting distributions. We now formally define the distribution over states encountered at the limit when an agent operates in an MDP under given policy, as well the conditions of existence of such a distribution. To that aim, we distinctly separate the discrete and the continuous cases, where the conditions to reach the stationarity of the system and the properties that hold at the limit of the interaction are different. In the following, we write M [s] for the MDP where we change the initial state şソ of M by s S.
(a) Finite state spaces. Let . be an MDP with finite state space and be a stationary policy for M. The measure & ASS) with (s') s) distribution giving the probability for the agent of being in each state of [[s] after exactly t steps.
x P. Suaey auno 8t s' {) the
Definition D.5 (BSCC). The subset E a strongly connected component (SCC) of M if for any pair of states s, E B, there is a t e N so that 5t (s' s) > 0. It is bottom SCC (BSCC) (i)
a maximal SCC, and ((i for each E B, P (B 8)
The unique stationary distribution of B is & A(B), defined as (s) Egng, Pa(8 8) limp-+x0 + Σi=o Sn [s] 81) for any 81 E B (Baier & Katoen 2008
The period of state E S is the greatest common divisor of the set {1 0 ( 8) 03.
A BSCC B is aperiodic iff all states of B have period one.
An MDP M is ergodic under the policy if the state space of MM consists of unique aperi- odic BSCC. In that case, & lim $t s) for all e S.
(b) General state spaces. For convenience, we focus on a subset of well-behaved Markov chains, namely Harris chains (e.g.. /Athreya & Ney] 1978 Revuz 1984), which are guaranteed to converge to a stationary distribution at the limit under some conditions. Let M be an MDP with complete, separable state space S. We denote by E(S) the set of measurable events of S. Let be a stationary policy for M. The measure & S A(S) with (5' s) P P.Mial(( Soze: au.. 8; E 5 }) is the distribution giving the probability for the agent of being in state s' E S after exactly t steps, starting from s S, for any S' e E(S). Denote the reachability event H e E(S) as
H (80:00 agse) e Traj there a N so that 8 e H}.
Definition D.6 (Recurrent Harris chain). The induced MC M is a recurrent Harris chain if there exist measurable set H E ESS), an e (0, and probability measure e A(S) so that
(i) И is recurrent, i.e., (0HH for all s S; and (ii) H a small set, i.e., for any s E H and C E E(S), (C s) > Ep(C).
Recurrent Harris chains are guaranteed have a stationary measure & A(S) Durrett 2010 Theorem 5.8.9), so that (S) Ei-g, P. (S) 8) for any S E (S))
Similarly to the discrete case, say that a Harris chain is aperiodic if and only if there is moment from which the agent is able to return to the region at any time: for any s H, there exists moment i N so that for any j > i, E (H s) 0.
18

---

Published as a conference paper at ICLR 2024
Assume M, is an aperiodic recurrent Harris chain with stationary distribution En. then S converges to &x in total variation as goes to o0 (Durrett 2010 Theorem 5.8.12), ie.,
lim drv (5%() s) = 1-
for all 8 E S.
To provide such a stationary distribution over histories, we define history unfolding MDP, where the state space keeps track the current history of P during the interaction. We then show that this history MDP is equivalent to P under म.
D.2 HISTORY UNFOLDING
Let us define the history unfolding MDP My. which consists of the tuple (Sx, A, Px. Rx *, *,7), where:
the state space consists of the set all the possible histories (i.e., sequence of actions and ob- servations) that can be encountered in P, i.e., Sx (A 22) hreset }. where we particularly denote by e the empty string with - (*()) = da, (no action nor observation has been perceived yet), and freset special reset state;
the transition function maps the current history to the belief space to infer the distribution over the next possible observations, i.c.,
E 8~7* of P(1s,,)) (11sss (s) Pu(6ar demmn a) de (h)) otherwise, where h of is the concatenation a, of with the history h (auT O resulting in the history (avrr, 0u.++1) so that ar = and 0T+1 o': and
p (h') h, a) P= h h,a) PP 8emaa a) Sr (h')
E
E Swavd (h')
if (+)) * Samen and
the reward function maps the history to the belief space as well, which enables to infer the expected rewards obtained in the states over the this belief, i.e., Rx,(,, (I-)+(1)) R(8, a).
We now aim at showing that, under the latent policy म. the POMDP and the MDP My are equivalent. More formally, we are looking for an equivalence relation between two probabilistic models, so that the latter induce the same behaviors, or in other words, the same expected return. We formalize this equivalence relation as stochastic bisimulation between Ms (that we know being an MDP formulation PP) and Mx
Definition D.7 (Bisimulation). Let M (S, 81, be an MDP. A stochastic bisimulation on M is a behavioral equivalence between states $1, 82 S so that, = 82 iff
1. Ra11, a) = R(82, a). and
2. P(T $1, a) (T a).
for each action E A and equivalence class T S/=.
Properties of bisimulation include trajectory equivalence and the equality of their optimal expected return (Larsen & Skou 1989 Givan al. 2003). The relation be extended to compare two MDPs by considering the disjoint union of their state space.
Lemma D.8. Let P be the POMDP of Lemma D.1 and म: B A(A) be a latent policy condi- tioned on the beliefs of a latent space model of Define the stationary policy म Sn A(A)
for Mn n a (. h)) a(- (h)), and the policy & for Ms encoded by the Mealy machine detailed in 2 Then, M and M are in stochastic bisimulation.
,7) so that Proof. First, note that the MC M is defined as the tuple (B x B, Roo, (br,br) b.),))).)
P 1b b,b)
E ( b.b.,.b)) P& b.a)
(1=(116)
19

---

Published as a conference paper at ICLR 2024
E a~x(-16) s-b --bbeePP(-s..e -(12,,))))) (b(boaa) (8) (b,0,, (6), and
E
E
Rgo (b, b) = E a~F(116)
E R(s
(cf. Definition D.4
Define the relation =y the set { (h, (b, b)) SN (Bx B) (*)))) b and 40* (h) = b). We show that = is bisimulation relation between the states of M and M Let h (Cor--1, Our)) Su. be B, and b E B that OT # o² and h = (b. b):
Rg (h) (a-x(1p*()) (e)) R(s, a) -a- (6) Es-b R(s, (b, b).
Note that histories are deterministically mapped to a unique belief (resp. latent belief) state through the belief update function (resp. ş). However, and & histories can lead to the same belief (resp. latent belief) state. Then, each equivalence class T can be associated to single belief and latent belief pair, aggregating multiple histories that lead to the same pair, through and sp Concretely, let 6' B, B e B, an equivalence class of = has the form T (b,, I
injective, multiple
so that
(a) the intersection of [66 with Sx is the set E Sx (1h') b' and x*((h) / }.
and
(b) the intersection of [(,,]] with the state space of M merely consists of the pair (V,8'>
Therefore,
P= (1V, 1h)
E
E
E P()
E (((35a)
dn (h') dh'
RB ((0151912 am=(( w (h)) s
1 a~a(1110 (N)) 9~ (M) s' p()a,a) o-
E
E
E
B
Snad (h1) 5, 5n (x)(N(())) (b) 8, (x) an
(by definition of [b [b,,
E
E
E
E
(r+(har))(b) - (m
s~+*(h)) x' -' ss.aa) of ((ls',a) (.vF(1l+(h)))
E
E
E
E
87(h)a00) (b') 8 ((.0000))00
(since = (b, (b,b))
a- (-1b)
((leeaa)
Poo (b, bb
Po (105 ((0,,8>)- b.b
By 1 and 2 we have that My and MM are in bisimulation under the equivalence relation = when the policies म and 7 are respectively executed in the two models.
Corollary D.9. The agent behaviors, formulated through the expected return, that are obtained by executing the policies respectively in the two models are the same: om.) EEMM [22= Ro(bb,
EMM [R(s,, do) + -1 Rx(0o.-1,
Proof. Follows directly from (Larsen & Skou 1989 Givan et al. 2003): the bisimulation relation
implies the equality of the maximum expected return in the two models, where the maximum taken over the set of all stationary policies of the two models. Since we consider MCs and not
MDPs, the models are purely stochastic, so there is no nondeterministic choice linked to the choice of action, and then there is only one possible expected return, which yields the result.
O
Note that we omitted the super script of in the main text: we directly considered i as a policy conditioned over histories, by using the exact same definition.
20

---

Published as a conference paper at ICLR 2024
D.3 EXISTENCE OF A STATIONARY DISTRIBUTION OVER HISTORIES
Now that we have proven that the history unfolding is equivalent to the belief MDP, we now have all the ingredients prove Lemma D.I
Proof. By definition of Mx. the execution of म is guaranteed to remain episodic process.
(a) Assume that the number of branchings the history space is finite when unrolling all possible episodes of the environment when म* is executed. Then, M is finite as well. Every episodic process in MDP with countable state spaces is ergodic (Huang) 2020), there is thus a unique stationary distribution Hz* lim S +) defined over the state space of Mx
(b) More generally, assume now that 12 is complete and separable space. We first show that M defines recurrent Harris chain. Take H (rreset intuitively the event of visiting the reset state is triggered infinitely often with probability one under any policy, which ensures the history unfolding to be Harris recurrent. To see why, notice first that by Assumption 2 MH is episodic, meaning that under any policy, the environment almost surely reaches rreset (and so, hreset by definition of Mx), thus condition (i) of Definition [D6 is fulfilled. Second, since H merely consists of a singleton, we can take p E A(SH) so that, for any measurable CE E E(Sx). p(C) = En- (- (:1Rro)): Pxe(C) /treset, a). Then, condition iii) of Definition D.6
is trivially satisfied. Therefore, MW " recurrent Harris chain, we know there exists stationary measure Hg over the histories of P that are likely to be seen under x Durrett 2010, Theorem 5.8.9). Then, the e-perturbation trick of (|HHuangg222)eerrr be used as is to enforce the aperiodicity in My So, we have that < converges to H. in total variation as t goes to oc (Durrett 2010 Theorem 1.).
The resulting stationary distribution is thus the limiting distribution over the histories of P when the latter operates under म, or equivalently, the limiting distribution of the MC M by Corollary D.9
E DISCRETE LATENT VARIABLES AND TEMPERATURES
As mentioned in Section 2.3, the optimization process of the WAE-MDP relies on temperature parameter, À E (0, 1). The latter controls the continuity of the latent space learned. The purpose of the parameter is primarily to learn a discrete latent space model: precisely, we use continuous relaxation of discrete random variables (Maddison et al. 2017). This is essentially the Bernoulli version of the Gumbel softmax trick (Jang et al. 2017p. The technique yields re- re-parameterizable and convex densities, which is compliant with stochastic gradient descent. The dero-temperature limit (i.e., passing from the continuous to the discrete setting) is theoretically enabled via simple rounding of continuous random variables. Furthermore, the logits of discrete densities are guaranteed to be identical those of their relaxed counterparts.
Alternatively, one could use the straight-through gradients estimator (Bengio et al. 2013). as used by van den Oord et al. (2017); Fajtl al. (2020) and Hafner al.) (2021). This consists in using discrete variables and non-difierentiable functions in the forward pass of the input through the neural networks, while continuous variables and surrogate functions are used in the backward pass, i.e., during the backpropagation of the gradients. This yields low-variance, but biased gradients. In contrast, continuous relaxations allow interpolate between these phenomena: higher temperature produces low-variance but biased gradients, while lowering the temperature to zero increases the variance but ends up producing unbiased gradients.
Therefore, at any time, the discrete densities (zero- tero-teemeeatture limit) of the WAE-MDP are used, except when the gradients of the objective are computed, where their relaxed counterpart are used (X > 0). As such, the WBU training procedure follows the same principle as WAE-MDPs: using continuous relaxation of the discrete random variables (precisely, multivariate Bernoullis) when the belief loss (Eq.8) ()) is minimized while the actual discrete (autoregressive) density is used otherwise. We chose temperature values by following the guidelines from the original paper (Maddison et al.) 2017). In the latter, it is mentioned that setting up annealing schemes (to the zero temperature limit)
21

---

Published as a conference paper at ICLR 2024
is a good practice but is not necessary for obtaining good results, which is confirmed experimentally in the experimental evaluation WAE-MDPs (Delgrange et al. 2023 Appendix B.8).
F VALUE DIFFERENCE BOUNDS
This section is dedicated to proving Theorems 3.2 and 3.3 Both Theorems bound the value differ- ence of histories, in the original and latent space models via our local and belief losses, to provide model and representation quality guarantees. Before proving the Theorems, we first formally define the value function of any POMDP, and then illustrate intuitively the meaning of each loss used to bound the value differences.
F.1 VALUE FUNCTIONS
We start by formally defining the value function of any MDP.
Definition F1 (Value function). Let M (S,A, R, 81,7) be MDP, and be policy for M. Write M[s] for the MDP obtained by replacing 81 by E S. Then, the value of the state 8 S is defined the expected return obtained from that state by running æ, i.e. (s(s) = EMill [>i=o R(s,, [. Let M = (Sa: PT. Rr, 7) be the Markov Chain in- duced by (cf. Definition D.4. Then, the value function can be defined the unique solution of the Bellman's 's equation Puterman 1994y: V=(s) (*(s) + Ep'- by V V=())). The typi- cal goal of an RL agent is to learn poliey x² that maximizes the value of the initial state of M: max.* (s1).
Property F.2 (POMDP values). obtain the value function of any POMDP P (M, O) by considering the values obtained in its belief MDP Mg (B, A, Pu. Ro, RR,bi,77 Therefore, the value of any history h E (A- -2)* is obtained by mapping h to the belief space: let a be policy conditioned on the beliefs P,, then we write V (h)))orr (= (77((1)) (h)). Therefore, we huve in particular
for any latent policy =: B A(A):
(h(
E TO
MMs++(NN) (NN]
r Re(be, ae)
(ef. Lemma D.8 for definitions of n and क*)
8 Σ (Ru((bes
E MxlM| [h]
(cf. Corollary D.9
n7
1=0
E [ Va (1)1
Ru +
(by Definition F1
a- (-1b)
N-Pu((Mh)))
E a (b))
Rx(hi, +
E N-Pa((ha)
[7 [V-V=(10]
(by definition of π
[ V2 (h. a...a o]]] (by definition of M.;)
E
R(s,a)+ + s-P((aa) o'
a~x(+le**))
-~
Similarly, we write V= for the values of a latent POMDP
F.2
LOCAL AND BELIEF LOSSES
Theorems
3.2 and 3.3 involve the minimization of local (LR, Lp, Lo) and belief (Lr, pi L losses. We intuitively describe how these losses are minimized via the latent flows depicted in Fig.
6
The procedure allowing to minimize the local losses depicted Fig. 6a At each step, state s, an observation of 8, and an action are drawn from the distribution of experiences encountered Hi while executing न. First, (s, o) is mapped to the latent space via the state embedding function of the WAE-MDP: d(s. o) 5. Then, the action is executed both in the original and latent space models (respectively from (s, o) and s), which allows to quantify the distance between the next reward and transition produced in the models. Finally, the original model transitions to the next
22

---

Published as a conference paper at ICLR 2024
Ma
Hm
M2
M
s
S
5t
h = (1 Oz:e)
B
bt
(l
a
6 x*(h)
R
R
L
P
P
DKL
St
5
5
T" (h)
R
LO
RR
P
r 1
P
(a) Optimization of the latent space model pa- rameters (i.e., R. P. and O) by minimizing local
M
S
losses.
(b) Optimization of the belief encoder y by min- imizing the (proxy) belief loss, as well as the re- ward and transition regularizers.
Figure 6: Latent flows used to compute the local and belief losses. Arrows represent (stochastic) mappings, the original state-observation (resp. latent state) space is spread along the blue (resp- green) area, and the latent belief space is spread along the yellow area. Distances (and discrepancies) are depicted in red. Notice that the blue area corresponds to the state-observation space Sn, which accessible during training (Assumption
state-observation pair (s', o'). Embedding this pair again to the latent space via d(s', 5' allows to quantify the distance between the original observation of and the one that is produced in the latent space model, from g' via o' O(- 8²). Notice that in practice, O is not necessarily Dirac (cf. Remark 3)
The procedure allowing to minimize the belief losses is depicted in Fig. 6b This time, the opti- mization performed on-policy, which means that is performed while executing the policy in the original environment. At time step > 1, the current history ends up in the observation o of state 8t. First, the discrepancy between the latent belief obtained via our belief encoder b, 40 (h) and the obtained via the true latent belief update function b' 7 (h) is evaluated and can further be minimized from that point. Second, we compute the reward and transition regularizers by minimiz- ing the distance between rewards and transitions produced from believed states 5, ~ be (i.e., states witnessed by the current belief bx) and those produced by mapping the current state-observation pair into the latent space, via (fs 0,) s, when the action . b.) is produced. Finally, is executed in the original environment and the process is repeated until the end of the episode.
F.3 WARM UP: SOME WASSERSTEIN PROPERTIES
In the following, we elaborate on properties and definitions related to the Wasserstein metrics that will be useful to prove the main claims. In particular, Wasserstein can be reformulated the max- imum mean discrepancy of 1-Lipschitz functions. The claim holds when the temperature is at the zero-limit, i.e., when the discrete variables and densities of the models This further makes the distance metric associated with the latent state space go to the discrete metric 1 X {1,0} (Deigrange al., 2023), formally defined 1, (x11,
used (cf. Appendix E}
iff x1 & X2- Definition F3 (Lipschitz continuity). Let X,Y be two measurable set and f: X - be function mapping elements from X' to y. If otherwise specified, we consider that f is real-valued function, i.e., y R. Assume that X is equipped with metric d: X [0. .0). Then, given constant K > 0, we say that ff is K K-Lipschitz iff. for any I, D2 E X. (f(x)) - f(xz)] & d(z,, Ø2). We write FK for the set of K-I Kppischhitz functions.
23

---

Published as a conference paper at ICLR 2024
Definition F.4 (Wasserstein dual). The mnntorovich-Rubinsserr duality allows formulating the Wasserstein distance between P and Q as (P(..Q) Supperj [R f(z) - Ey-q f(y)l.
Villlnn 2009
Property E5 (Lipschitz constant and total variation constants). Let f: X R, so that is a metric on X. Assume that f is K K-Lipschitz, Lee, f E F5. then for any distributions P,Q (
[Ea-p fxx1) Eaz-Q (x22) & K Waa (P,Q). In the case where d = 1,, for any bounded function g: X Y with Y E R. we have [Eni g(21) (x.=00((>)) & Ky W₁ (P,Q)
Ky Ky-drv(PP Q), where Ky 2 suPEEX (g(x)) (uülerr 1997 Sriperumbudur et al. 2009)
In both cases, we call the smallest K (resp. Ky) the integral probability metric (IPM) constant of the function f (resp. g).
Property R5 intuitively implies the emergence the Ky constant in the Theorem's 's inequality: we know that the latent value function is bounded by supr,a (R..))/1 so given two distributions P, Q over S, the mean discrepancy of the latent value function is bounded by Ky (P,P,QQ in the zero-temperature limit.
Finally, since the value difference is computed in expectation, we introduce the following useful property:
Lemma F.6 (Wasserstein in expectation). For any function f: y x X IR so that X is equipped with the metric d, consider the function 9y X - R defined gy
fly x). Assume that for any y E y, 9y is Lipschitz (or simply bounded if d 14) and let K be the IPM constant of 9y- Then, let D A(Y) be distribution over y and P. A(X) be two distributions over X, have Ey Bay-P [y, x1) - Eaz-Q 22) < K Wa (P,Q).
Proof. The proof is straightforward by construction of gy:
E
f(y - E f(y,22) -1>P
E
x2~Q
E n~D
E
(y(x))
-
E (1(222)
-1~P
x2~Q
<
E
[K Wa (P,())
(by Property F55
n-D
K (a()QQ)
F.4 MODEL QUALITY BOUND: TIME TO RAISE YOUR EXPECTATIONS
Let us restate Theorem
Theorem 3.2 (Model quality). For any latent policy n: B - A(A), the values of and P are bounded by the local and belief losses in average when, in both models, the actions produced by म, which is conditioned on the latent belief induced by s, i.e., a ~ #( ( (h)):
L
E
- V=(()) < 1-7
(6)
h~
Proof. The plan of the proof is follows:
1. We exploit the fact that the value function can be defined the fixed point of the Bellman's equations;
2. We repeatedly apply the triangular and the Jensen's 's inequalities to end up with inequalities which reveal mean discrepancies for either rewards or value functions;
3. We exploit the fact that the WAE-MDP is in the zero-tenoe troperrrrrrtrre limit to bound thos discrepan- cies by Wasserstein (see Porperty F5 and the related discussion);
24

---

Published as a conference paper at ICLR 2024
4. The last two points allow highlighting the L1 norm and Wasserstein terms in the local and belief losses;
5. Finally, we set up the inequalities to obtain discounted next value difference term, and we exploit the stationary property of Ha to fall back on the original, discounted, absolute value difference term;
Putting all together, we end up with an inequality only composed of constants, multiplied by losses that we aim at minimizing.
Concretely, the absolute value difference can be bounded by:
E Mm
(h) - V (h)l
E Mm
E
E
()s,a)+7
E s/ o' O(1o',)
E
vi(hha.dd)
g~ (h) a~ (-Ip*(())
E
E
R(5,a) +7
E
E
vz(h.a.o))
5- (l) a~ a~ ( (N))
s (1
(see Property F2)
E
E
R(s, (((s,)++ +
s'-P(-)s,) d' as
E
E
h- hn (1)))
( (h)
E (1b)
R(5,a) +7
E
E
V Vx(h.a.0)
5' (- 0
(Jensen's 's inequality)
E h~ H= a~ (+)))((())))
E
R(
R(5,a)
8~F+(h)
V (b
V (h gw)+(h) (-((s) (Triangular inequality)
?
p (h) s P(ls.)) o' (c(ora))
For the sake of clarity, we split the inequality into two parts.
Part 1: Reward bounds
E hs a
E (p+)
E + (N)
R(s,a)
a)
E R(s,a) a)
3-17(()
E
E
E (>)*(()
[R(s,a)
a)]
E (R(}(s,o),a) R R(,,a]]
1,o a- (b))
eNT (h) (h)
(o the last observation of h; the state embedding function é that links the original and latent state spaces comes into play)
E
E (b)
IE (N)
[R( R()(s,
E (A) (1)
[R(o(s, o), a) - R a)]
No0 Hv am F( (11+
( (Triangular inequality)
E
E
o))
[R(o(s,o), R R a)]
h,0- (xF((1(*())))) (b))
(k)
1 (N) 5- (h)
(Jensen's 's incquality)
=
E
E
E
R((s,)) a) - -((o(
hoo Hs a (1)))()))
(
+
E
R( o). R(5.a)
h,o (~T((1()(()))
SeTT (+))
25

---

Published as a conference paper at ICLR 2024
= LR E
E
E
E
[ - R(5,a)] (by definition of LR)
ho a> (1))
| (h) (1))(())
=IR +
E
E
E -~ (h) (1) S (h)
E
E
[[R( o), R() + [R(8 - R(s.a]]]
h,0 a (-1p (b))
( (the belief encoder yp comes into play)
=LR +
E
E
I
E
[R(((,,0), - R(5,a)]
h,0 (12()()())))
p~ (h) (m)
E
E
(R(s1, - R(s
8~7 (h) AT (R)
LLR +
E
E
E
E
(R(((s - R[3,,)]
1,0 a (110*()))
8--7 (h) (
E
E
R((51,)) - a)]
a (1)
(Triangular inequality)
EIRR
E
E
E
R()(s,0). - -R(s,a)
ho0 (-*((11* (b))
swi (h) snpp((h (h)
E [R(51,0) a) R(5,a)) a)
5 (1))
(Jensen's 's inequality)
=LR
E
E
R( R) |
ho (nF(((+(()) s (
(h)
+
E
E
E
E s1
R(51,a) -R(5a))
h,o
( (N))
=LR +
+
E
E
E (L) (N)
E
[R(5, a) - R a)]
h,0 a~ (-1p((()
(by definition of Eq.
<IR
1 + E
R W'
( (as A 0. by Lem.
F.6 and Prop. H5
R
h~lr
=LR + R + R Lo:
Part Next value bounds
7 E h
E (N))
E
E
E
V V2
of
V va(h-aad))
E
E
E
(M)
o' (118)
F
E (N))
E
E
E
E
(z ah..a.07 a. o")
H- a~ a1 (-1(
8 P())) 0'
E
E
E
E
(=
s' of of 8 000 SNT (b)
a
3 (N) x' P( 0-13)
(the state embedding function & comes into play, as well as the latent observation function i)
26

---

Published as a conference paper at ICLR 2024
?
E
E
E (a) s' a.0) of (-)orsa)
V ah.h.a.00)
E 0-0(1o(s,0))
ax(
ho H a- (1))+())))
+ 7.
E
E
E 80
V ah
He (-)()()()))
s~**
V (a (1E00
(Triangular inequality)
7
E
E
V (2h (~0((16(),01)
E
V (a(
h- h- (-1l0+))))))
(1)) s'~P(-la..) d' O-
+ + E
E
E
E
vz(h.a.0)
ho a~7(11100000 (-)(+10*())))
(h)
st P.dd-Pa(-ss..a 8' (1D
E
V aa
((dd(()))))))) (110
E
E
E
V Va(h.a.o)
B (b) & P (1)0 a) of 0(18)
V (h
E
E
8~T (h)
8P( x
(o is the last observation of h; the latent MDP dynamics, modeled by P, come into play)
?
F h- Mr a~
E (-) (110
Vj Va (~0(11())0)
E
V ah.
s(-P(-)).0) o' (11osa)
+ 7
E
E
E
E vz(h. 00))
h,o H. a-a a-)(1-1+(())
s-)+((h)
s.d-Pa(t-sa..a) (x (-
V Vz(h-a.6))
E
+ 7
E E
E (h) o - (((s)
V al
ho a (h))
E
E
E
V v a. of)
&- (1) g (( o' (113))
(Triangular inequality)
&
7 E
E
E
E
E
Va (ho
(-0(16(6 - va(hea.00)
E
l- He a~ (h))
( (())
p-P o' (-1esaa)
+ Y
E
E
E
E
E va(h-a-o)
hoo a>T((>p*()))
a~ (b)
R pa( d (((88)
v=(h-a.d)
E
E
+7. E
E
E
(h) s' P (.(5)) (=(
E
E
(h)
27

---

Published as a conference paper at ICLR 2024
E *-T*()) (N) of -f (by definition of pp and the Jensen's 's inequality)
V (h a dd)
7
E
1/2 ah. (-D(1o(s)0)) (k) of P' x. o' o'
E
V (a(h-a-0 ath.a..)
h~ Hw a-)(-1++(h))
(h)
+ 7
E.
E
E
K W3 (pPal 8,0), ((.1((8,0),00
ho a- (1))))))))())))))))))))))))))( )n)*(()) (h)
+
E
E
E
E
E
V (h
Ho
swuT 5' P()p(())))))) 0113)
ax(h-a-d)
E
E
E
((1 d' (-118)
(as - 0, by Lem. F6
9
E
E
E
E
E
Va (a -
E (-D(1+(((0))
V aa(
He ama (
(b) s'-P(-)).) 0'
+ y KK Lp
+ +7
E
E
E (h) 8' ( d'((5))))
E
V V ah))
He (-10+(())))
I
E
E
V (h. a.)
8 (1
of
(by definition of Lp)
E
E
E
E
E
V Va(h-a.0)-
E
V Va(h-a-6')
h~ am (1))
(h) s' P( O1116,a
+ K Lp
Va(h-a.o)
E +7 ho (-))+(())))))
E
E
E
E
* (1) x' -11
Pa(h-a.d)
E
E
E
8 p (b) s' P (-1a,a) o' ((K)
+
E
E
E
V ah.
E (-)*(h) (- o'-5)
E
E
(a
gup (b) s'- P. (( (1)5)
(the belief encoder & comes into play)
<
E
E
E (h) P a. Ot-16,,
E
E
V (h- - --3(16()00)))
E
aa(
h- a (())
+ Lp
+ 7
E No
E
E s (l) x (( (113)
E
V=(h-a-o))
a-al
vz(h-a.d)
E
E
E
8 3 (h) of (.lsa))
(1)8))))))))))
+7 E
E N- a~ (1))
E
E
E ah(h-a.d)
s (1) 8' P (( of -
28

---

Published as a conference paper at ICLR 2024
Vz aolh.a.d) a. d) (-)s,) of (-100)
E
E
E
8-77
( (triangular inequality)
7
E
V/2 (2
E
V (a
h~ Ha (-)(-1)+(())) (1) s' (,),)) of (+1ss,))
(-0(1o(s
+ KKy Lp
+ +
E
E
E
I
E
R
ax
ho (-)(- 8 (h) (1))
8'~ ((1000)))) (110)
E
E 0'~0(155)
v=(h.a.o)
PPP 8.0)
+ 7 E
E
E
E
E
Vx ax(h-a.o)
h~H (-1(+())))
(h)p(()) of PPPP(s of (-100
Vz(h- aahh.a-0 a. d)
E
E
E
8-
(-15,) of o' (-18)
(Jensen's 's inequality)
7
E
V Vh
E --3(((()
V a
h- H~ a.llll (h)) (1) s' 0.0) o' O-1ss
+ < Lp
+ 7
E
E
E
E
K₂ Wa (P(. (l),),, |(.1(,))
ho a (-1x++
em (1)) (N)
+ 7 h- H((p+())))))))
E
E
væ(h-a-o)
( (h) 3-P(1) of (-13))
V V (h.ha..)) 8~ (b) H' Pl (+18)
( \ 0, by Lem. F6)
- E
E
E (1) s'-P(-)))) o' (c1or,a)
(ahh.a.0))
E (~D(16(().0))
D ah.
l- Ha a~ (1)))+()))
+ ~ky (Lp L1)
Vi(h-a.d)
+ 7 E h- H. an (-))+((()))
E (h))
E
(-l,a) of (+()")
V vath-a.0)
E
E
E
5 5-
(1 o 0+188)
(by definition of Lp Eq. 5p
7 E
E l- u~ a~F(-1(*(h))
E
V/2 Va
E --3(11(),00))
V (1
(1)) s' a..0) (-16)-
+ 7 K (Ip 1.))
+ 4
E
KrWa (7)
h~ A~ a= (k))
(as \ 0, by Lem. F.6 note that Wasserstein symmetric since it is distance metric (Villani 2009
29

---

Published as a conference paper at ICLR 2024
7
E l- W (-)()))*(()))) (o)T((h) s() of oH1osa))
E
E
V7 (a.
E
V Vx
0 2~0(16(5,00)
+ K (Lp Lip + Lrz)
(by definition of L,,, Eq. 5
(Va(h.a.d) V ax(h....d))
= 7
E
E
E
E
h,owH, (1)++(()(()))
+ vv (~D(((()
Va 01)
+ +k (Ir+ + 1i + l1z)
7
E
E
E
E
[V -V
h,0 Hs (-
o' Pattlla.)))))
Va(h-a.o)-
Vi(h-a-d))
+ +7 ho a (-)))))*())))
E
E
E
E
E o d'D((
(h) o' Pu((
kT : (pp + 1p + ltr) TTr)
(triangular inequality)
[v<(ha.o)) - -1 (h of]
& 7'
E
E
E
E
Ho a~ (-T(++++)()))))
(h) 8' PatllPa
+ 7'
E
E
E
E
Va(h-a.o') 8 (lo(ee)))))
E
vhh
ho (-x((
a~ (b) p() 01-0(-1s,a)
+ k
(pp
ip- + L17
(by definition of Pn and the Jensen's 's inequality)
[V - V (h. a. .o]]
<
7
E
E
E
E
h Hn a
(l)) ( Pntt-la..)))
+ E
he a~a a-T((- B
E
E
E
Kydrv (0(-188,0)
E d's''a
0(11((,0)))
(~)*(h) s p(((aa
+ <K (pp Ig + Tup
(cf. Prop. F5 and Lem H6
= 7. =Y.
E h,o Ha (x)()
EE
E (d
E
Va V t ah
+ 4y (Lp 4p + Lo + Lo)
(by definition of Lo. Eq. 4
7 E
E Ho a~ (>T(-1o+(())) + (N) of ( Pa(.)a
E
E
(z(ha..d) - V (h. a. -o)
+ y KT , (1p + T.p + Tor + T.o)
(Jensen's inequality)
7- E. (V+(h) - V v=(h)] + >Kg (pp + LI + LT L Lo) h,
(H₂ is a stationary distribution (Lem. D.I which allows us to apply the stationary property)
Putting all together. To recap,. by Part I and 2, we have:
EE h- Hm
[V>(h) - V (h)) < LR L + R Li +7. E VV((h) ((())
h~7t
+ kk (Lp Lip + L Lo)
E. (((1)->(hh))))- - x LRR 1 + R L= + 7Ky (Ip Li L + Lo)
H4
IRR + _ R' L + kKp
(Tr + Tiğ + I + 1.o)
E h~
VG(A)- - v &
1 1-7
30

---

Published as a conference paper at ICLR 2024
which finally concludes the proof.
F,5 REPRESENTATION QUALITY BOUND
We start by showing that the optimal latent value function is almost Lipschitz continuous in the latent belief space. Coupled with Theorem 3.2, this result allows to show that whenever any histories are encoded to close representations, their values (i.e, the return obtained from that history points) are guaranteed to be close as well whenever the losses introduced Sec. (32 are minimized and go to zero. Phrased differently, this Theorem ensures that the representation induced by our encoder is suitable to optimize the value function since the distance between beliefs in the latent space characterizes the distance of behaviors of the agent in the original environment. The latent belief space thus captures the necessary information to learn a policy that optimizes the expected return.
Definition E.7 (Almost Lipschitzness), Let be a measurable set equipped with a metric d: X - [0, 20) and f: X - R. We say that f is almost Lipschitz continuous (c.g.. Vanderbei 1991) iff for all e > 0, there is a constant K 3 0 so that (x1) - f(x22) < Kd(z1, x22 + for any T1, D2 E X. Notation (Optimal value function). For any MDP M, let n' be an optimal policy of M. then we write V* for VT (see Property [.2 for the values of a POMDP).
Lemma F.8. Let P (M, S2 оy be a POMDP with underlying MDP M = A.P. 7). Assume that P is finite, i.e. S, A, and are finite sets. Then, V* is almost Lipschitz continuous.
Proof. Define as the set of real-valued bounded functions V: B R and U: B AxV as
U(b,a, a, V) Rs(b, a) + E M~Pa(-1b,a)
[[)]
The Bellman update operator is defined as 2t V V as (LV) = max.5A (b, a, V) and an isotone mapping that is a contraction under the supremum norm with fixed point V*, i.c., V - =UV (Puterman 1994, [auskrecht]] 2000, [Sutton & Barto 1998). Furthermore, for any initial value function Vo E V. the sequence resulting from value iteration (VI) V++1 = UV, converges to V (with linear convergence rate PPuterman (1994): for any c' > O, there is a i N so that for all j > i, V & é'. Now, let e > 0: in particular, the latter statement holds for = 9/2. Since the convergence of holds for any initial value, we assume that Vo E V has been chosen as piecewise linear conver (PWI.C) function. Then, for all > 0. V/ is also PWICC (Sondik) 1978 Smallwood & Sondik 1973 Hauskrecht 2000). This means that V, can be defined a finite collection of linear functions. Thus, V, is also 2K 2i-Liischitz: one just need take 2K as the higher slope of these functions (in absolute value). In consequence, for any pair of beliefs bj. b2 E B.
[b,,,, - V*bb2)
= =IV"(bi) - vi(bi) + v2(6r) V,(62) + V((b2) (b2)
< <\V'(b,) - v((br)] + +VV((b) v((6)) + +VV,(b2) (b2)
(Triangular inequality) (by the convergence of VT)
<2c + [V,(bb) - -V2(22]
= + [V,(b,) V(((2)
< + K. drv b2),
(since drv b2) = -//2/0 - ball1)
which means that V is almost Lipschitz, by definition.
Lemma F.8
be applied to P. To see why, notice first the state space of P is finite. Second, O is deterministic by definition (it consists of the Dirac centered on the observation produced by O,), so the observations of P can be limited to the union of the supports of (- $), ranging over all latent states & 5. The resulting set is also finite. Note that the same reasoning holds for any latent observation function with finite support. Therefore, we obtain the following corollary.
Corollary F.9. When the WAE-MDP in the zero-temperature limit (see Appendix E. the optimal latent value function of P is almost tipschitz-onntimnnn
Theorem 3.3 (Representation quality). Let x= be an optimal policy of P, then for any > 0 and 3 0, there is K 0 so that for any histories hy, h2 of length at most that are measurable
31

---

Published as a conference paper at ICLR 2024
under P. and P with (h1) = b1 and 40 (h2) b2. the representation induced by yp yields:
KL + L
1
+ 1
Hax ( (7)
[Vz. - V7. (h2)
KWa (b1, + E
1-7
(kx)
Proof. First, observe that for any measurable history h. (h) - x. & (1)
BA (k') (h') (see. e.g.. Gelada et al. 2019
Let > 0, n 30 0, and h1 h2 be two measurable histories so that h; (Aur-1,00.) 0 with T & n and Ot E (- s)) for some $e e 5 and for all < T. i {1,2} (the two histories are compliant with P and P). Then, we have:
V (h1) Vaa (h2))
= -| (h1) - + + V x= (h1) - (h2) + (h2) - V (h2)
< V/ (h1) - V= (h1) + V (h1) - V (h2) + (h2)
(h2) - V (h2) (Triangular inequality)
<Hz (h1) Na
[/+ (h) [1) + V (h1) V (h2)
+ Ho. (h2) (h2)
E
11x (h) - v.
h-H,,
< (11) - Vz (h2) + 1- - (Hx. -1 + Hax
L
(Thm. 3.2
Consider h1 arg min { (ggg (h): h E supp(Ha** n H(P, PP) and |/| & n), where supp(H=.) denotes the support of H.. |h] the length of the history h. and (P,PP) the set of histories com- pliant with both P and P, i.c., histories whose observations belong to USes supp(O( that h1 exists: we consider histories which are measurable in the two models (i.e., the original and latent POMDPs) and the observation space of the latent POMDP is finite (cf. the premise of Corollary E.9. So, the set of histories from which we extract hy consists of a finite number of action-observation branchings.
(3)) Notice
Recall that the latent value function (defined over the latent belief space) is almost Lipschitz con- tinuous (Corollary (.9). In particular, for d (14))))( there is a K 3 0 so that for any
6, b' E B. (b) V-))) & kWa(k, + 8. Then:
V*(hi) - (h22)
=IV"(F*(hii) - (*(F((h2))
- - (((") ((66)))))) VM(
- +1(6(())) ("( +
- (
(Triangular inequality)
(H+
17 - + (4)6 [-bb)l
h~H
+ +Ha (112) -1
F
['(++(1)) (*())(h))
1- Hy.
(1 -1 + H1.x (h2) (bz)-1)
E
['(++(h) - + - (b(b2)
< (z. Hx. (h2)) 11)
E.
[KW] (7 (h), (()) + + + KW, (b,, 52) + +
K-H.
(V* is almost Lipschitz) (by definition of LT)
(Hm-(11) + ++x (12) -1 (KLT + +6++ + KW, (b,, b2) +
(zz (h1) -1 + 71 (h2) -1) KL + KW, b2), + 5 1 H1 (hz) -1 + H (122) (hz))-)
Har (h1) -1 + (h2) -1 KLG + KW, (2) + 5 (1 1 + 27 (h1))(1)) 7t)
32

---

Published as a conference paper at ICLR 2024
-(Hor- (h1) =(Hz=(hh)"' + Hz. (h2) KL + KWa (b1, +e
Putting all together, have that for any > 0, n 2 0, one can find constant K > 0 so that:
KL, + C (14x + H (h2) "
- (V (h2)l
KWa b2) +Et
1-y
Remark (Representation quality for latent policies). The bound of Theorem 3.3 also holds for any latent policy मे which is executed in both models by processing histories through ş. To see why, notice that such policies can be encoded through a Mealy machine (Definition(D.3) with observations as inputs (instead of states): one just need to use as memory states the reachable fragment of the latent belief space. Then, each time an observation 0++1 is perceived subsequent to an action ar, the memory b, of the machine is updated b1+1 yp (be, at Ot+1).
Remarkably, when a horizon is fixed, the set of latent beliefs that can be reached in the Mealy machine is finite since (i) P is finite, (ii) the initial latent belief fixed, and (iii) each belief has at most |A| | O(- 8) successors.
Each function V; obtained from a sequence (Vi),eN of value functions generated by VI to evaluate any policy encoded through a finite state controller is PWLC (Hansen 1998). Since VI converges, one can find a value function V; from this sequence (for a finite > 0) which ensures an error of at most E 0 w.r.t. the value of this particular policy. Consider i as being the horizon used in the memory of the Mealy machine. Then, all the arguments from the proofs of Lemma F8 and Theorem 3.3 can be applied to the latent policy encoded via this particular Mealy machine.
G ADDITIONAL DETAILS ON BACK PROPAGATION THROUGH TIME
In this section, we elaborate on why BPTT is not required, nor recommended, to learn representa- tion of the belief when using our algorithm.
In RNN-A2C (Figure?) 2, right) the hidden state is computed recursively, which is in the same spirit as how the belief computed through our WBU (Figure 2 left). The motivation behind activating BPTT for RNN-A2C and disabling it for WBU lies in the difference between beliefs and hidden states and what they represent.
On the one hand, beliefs state distributions which can be inferred from the previous belief and be themselves seen as Markovian states (in the belief MDP Ms). Therefore, one can see the belief update function the transition function of MDP, where observation is additionally provided. Disabling BPTT is thus consistent with the literature of model based RL for learning Markovian transition functions (e.g.. Gelada al. 2019 François-Lavet et al. 2019 Delgrange et al. 2022). Notably, our goal is not targeted towards learning directly a history representation (from sequences) but rather on learning to reproduce the belief update rule, based on the latent stochastic transition function P. and observation function O of the latent model. The latter are learned separately and considered fixed when learning the belief update (Section 4. Therefore. stopping the gradients does not produce approximate gradients. Furthermore, the error linked to this recursive update can be fully determined through the error of our latent model (Theorem 3.2) Then, enabling BPTT increases the risk of accumulating errors across time.
On the other hand, RNN hidden states are latent representations of the history and it is up to the RL algorithm to learn a useful representation of the history for the control task. This is in stark contrast to our approach: we separately learn representation which is guaranteed to be well suited for the RI. agent, and we use as input of the policy. In R-A2C the main objective is to find a policy that optimizes the return conditioned on the full history, making it sound to use BPTT in that case. By contrast, in WBU, the RL agent uses the belief as a given state-signal to condition the value function and the policy, and therefore, the RL objective is not backpropagated through the belief encoder.
Finally, it is worth mentioning that as beliefs are computed with a forward recurrence, errors may accumulate over time. This means that if the belief bi+b diverges significantly from its correct value, then b will diverge even more. Then, allowing gradients from the learning of b to be computed based incorrect +1 and further flow to b, may hinders the learning procedure.
33

---

Published as a conference paper at ICLR 2024
H ALGORITHM
We describe the final WBU learning procedure in Algorithm | Note that they keyword Update means that we compute the gradients of the input loss, and update the parameters of the neural networks of the pointed function/model accordingly.
Normalizing term. Given the set of parameters & of y, minimize the KL divergence DEL by gradient descent on the Monte-Carlo estimate of the divergence:
V. V,DKL (yp (be, ae, 04+1:4) ll (be, az, 0(+1))=
V.
E
log p (8(11 B at 01+1;)) - log E P(5 3,0;) - -log O S241) 5-
8441 (b.
Notice that the first term of the divergence (the belief normalization term) of Eq. 8 does not depend on & and thus yields zero gradient. Nevertheless, we observed during our experiments that adding the normalizing term allows stabilize and reduce the variance of the belief loss.
Optimizing Wasserstein. To optimize the Wasserstein term the belief losses, we follow the same learning procedure than (Delgrange et al. 2023 Appenix A.5): introduce neural networks Fa (for (.,)) 92 }) that trained to attain the supremum of the dual formulation of the Wasserstein distance. To do so, we need enforce the Lipschitzness of F and, as in WAE-MDPs (Delgrange) al. 2023), we do via the gradient penalty approach of Gulrajani al. (2017), leveraging that any differentiable function is I-Lipschitz notice that do not directly optimize the total variation distance of T.o but rather the Wasserstein; we take the usual Euclidean distance as metric over Ω which is proven to be Lipschitz equivalent distance converging the discrete metric the temperature of the WAE-MDP goes zero (Delgrange et al. 2023 Appendix A.6) to finally recover drv.
has gradients with norm at most everywhere. Finally,
I HYPERPARAMETERS
Table I provides the range of hyperparameters used in the search, along with the selected values for each environment. The hyperparameter search was performed using ÓPTUNA (Akiba et al.) (019). Pre-training of the WAE-MDP involved collecting 10240 transitions with random policy and per- forming 200 training steps. These pre-training transitions are taken into account in the reported results.
Additionally, (for R-A2C) and (for DVRL) present the specific hyperparameters used for each algorithm. A grid search was conducted over all possible combinations for both baselines. The hidden size of all neural networks was set to 128 neurons and two hidden layers (except for the sub-belief encoder, and the policy which uses one) without further tuming. The experiments were carried out using 16 parallel environments. The original implementation DVRL and their version of R-A2C were used in this study.
We ran the experiments on a cluster composed Intel Xeon Gold 6148 CPU.
In the following tables the enviromnents represented follows: STATELESSCARTPOLE V, SOUSYSTTATELESSCARTPOOOO + REPEATPREVIOUS Q SPACEINVADERS NOISYSPACEIN- VADERS
The t-SNE van der Maaten & Hinton 2008
presented in the main text, in Fig. was obtained using CuML. (Raschka al., (2020). Table 4 reports the hyperparameters that were used. We used 50,000 the training.
pairs, sampled from a dataset of I million timesteps collected at the end of
34

---

Published as a conference paper at ICLR 2024
Algorithm 1: WASSERSTEIN BELIEF UPDATER
Input:
Batch sizes BWBU, BWAB, BEXXT global learning steps N; no. of model updates per iteration : your favorite collect strategy Tint replay buffer (RB) D; Lipschitz networks FF: S
Fa Ω - I; observation variance network O.: and loss
weights W Wpp
collect =(8 8i, Os, as, Ts, s/, of li=1 by executing Timi for Nuir steps; store 0 in D
Mut
to Use the exploration policy to collect transitions and initialize the RB
repeat N times
repeat MoDEL times
Update the WAE-MDP model for NMODEL consecutive training steps
for i - 1 to BWAE do
s, 0₁ ag Ts. S,, d.) ~ - ( 03)
+ Sample a transition from the RB
t Embed (5 o,) to the latent space 1> Observe the resulting latent state via O LWAR compute the WAE-MDP loss transition batch 8g. Os, Ts, Ti, f, of )== Bwan Update the WAE-MDP components (in particular, those of Eq. 2 by minimizing LWAe Observation loss Update Fa by maximizing Lo and enforcing its 1-Lipschizness w.r.t. metric da
0, ~ (.
Lo 1/BWAE E.Ear [Fa(d.) FF [0)
Update O₂ by minimizing Lo
for i - 1 to BWBU do
Su - 81i 50 - 515 bo - 8za 30 + 31
to 31 is arbitrary, e.g., zerves
for 0 to I do
> Produce the action as according to the sub-belief B execute a in the environment, receive reward ft, and perceive the next state-observation (st+1: 0t+1)
store the transition (st, Or: de, Tt: St41, 9041,0++1) into D
3++1 - (sg((s ae, 0(+1)
+- Updatee the sub-belief: sg is stop gradients t> Retrieve the belief distribution bu+1 via the MAF WI
b++1 - -M(+))
Believe the next latent state E= Marginalize the next latent state distribution w.r.t. the current belief for j to BNEXT do 5 b.: C) flog [log P(5++1 5, (z) log BNEXT]
5t+1 -b,+1
c log bigbi+12 ($1+1) - LSE(1 C) log j=1
B
log O [(41)
Pointwise decomposition of Eq. divergence with the belief update rule
G R(o(s, or), az) - R(s,,,,))
t> Latent reward regularizer
P( $2, at)
- Transition to the next latent state from the current believed latent state
chi [F oc+1)) FF(5 P
LLttent transition regularizer
Update Fp by maximizing B.wau 21 cp
T-1 and enforcing its 1-Lipschitzness w.r.t.
latent metric d
Update T0 and M by minimizing l/(wwuu+ E IT-1 W L.
+ Wp
Update π by minimizing the A2C loss on the batch { (B d rd
roT Vi-1
Bwmi
function
(smooth) Observation Filter t> Decode 8; get the standard deviation of the reconstruction
(8)3) (8)
return (
35

---

Published as a conference paper at ICLR 2024
Table I: Range hyperparameter search and selection per environment.
Range
Updates per Beliel Update
I-2 leaky relu, elu leaky relu, tunh
Activation function
lenky relu tanh
leaky tanh
relu
clu
elu
relu
Activation function lipshitz
tanh
tanh
leaky relu
Activation CNN
leaky relu,
elu
elu
Policy conlig
Learning rate
Le-4, 3.c-4, 5.c-4 1. 10
3.c-
Le-4
1.c4
5,e-4
5,c-4
Clip norm
10
10
10
10
10
Bellef conlig
Same opt policy
Yes, No
Yes
Yes
Yes
No
No
Learning rate
Le-4, 5.c-4
1.e-4
3.c-
Loss factor
Lee-, Le-2, 1.e-3, 1.e.4, 1.e-5
Le-1
Le.4
1.e-5
Filter variance
1e-2, Le-3
Le-2
L.c:3
1e-2
1.c-2
Normalize log obs filter
False False
False
False
True
True
True
Sub belief prior temperature
0,666 .7,, 9,
0.66
0.9
0.66
0.99
0.99
Reward loss scale factor
0.1. I., 10. 50.,
20
Transition loss scale factor
0.1. 1., 10., 20., 50., 100.
100
50
50
100
Buffer size
4096, 8192, 16384
4096
4096
16384
4096
40%
Critic Update
10
N State samples
16, 32,
64
32
32
32
N next state sumples
12, 64
32
64
64
16
32
WAE config
Latent state size
0 10, Others: 18 25
5
18
Minimizer learming
L.e-3,3.e-,1.e-44,3 5.e., Le:5
5.c-5
5.c-5
3,e4
1ee55
1.e-5
Maximizer learning rate
3.-..3.c-4..e-44.55 5.c-5, Le-5
5.c-5
5.c-5
1.c-3
3.e-4
1.c4
State encoder temperature
0,33, 0.666, .75. 9,99
0.75
0.66
0.5
State prior temperature
0.33, 0.66,, .75.
0.75
0.75
0.5
0.75
0.75
Local transition loss scaling
10., 50,, 75., 80.
50
80
80
10
10
Steady state scaling
10., 25., 50., 75., 80.
80
25
N critic update
1,
10
Batch size
128, 128, 256
128
128
256
128
128
Clip grad
1, 10. 100.
100
100
100
10
10
Obs reconstruction weight
1.2
Obs reg. min learning rate
L.-3, 5.65, Le-5
5.0-5
Le-4
1.e:3
5.e:5
5.e-5
Obs reg. max learning rate
Le-3, 3.c-4,Le-4, 5.c-5, Le-5
5.0-5
3.e-4
5.e-5
1.e-5
3.c4
Obs reg. gradient penalty
50, 100, 500, 1000
1000
1000
Table 2: R-A2C hyperparameters
Range
Opitmizer Clip grad norm Learning rate
Adam, RMSProp
RMSProp
RMSProp
Adam 10. 5.e-4
RMSPrrp
RMSProp 10
0.5, 1., 10. 3.e-5, 1.e4,, 3. 5.e-4, 1.e-3
1.0
5.e-4
1e-4
1.e-3
1.e-3
Table 3: DVRL hyperparameters
Range
Optimizer Clip grad norm Learning rate
Adam, RMSProp
Adam
Adam
RMSProp
RMSProp
RMSProp
0.5, I., 10. 3.c-5, 1.e-4, 3.e-4, 5.e4, 1.e-3
10.0
1.0
1.0
10
0.5
5.c-4
Le-3
3.c-4
1.c-3
1.e-3
Encoding loss factor
1.11.5.05 5,10,15
0.5
1.0
1.0
1.0
0.05
Number of particles
10
5
Table 4: t-SNE hyperparameters
Method
FFT
Perplexity
10
N Neighbors
50
Metric
Manhattan
Learning rate
2400
Early exaggeration
50
Late exaggeration
3
Exaggeration iter
300
36

---

Published as a conference paper at ICLR 2024
J INDEX OF NOTATIONS
Divergences
DeL
KL divergence
drv
Total variation distance
Wa
Wasserstein distance
Extended POMDP
O¹ Sq 12 Extended observation deterministic function (projection the observation space)
P
Extended POMDP
Ma Extended MDP
Pa: Sq
(Sn) Extended probability transition
Ry: Sp A R
Extended reward function
Sp = 8 Ω
Extended state space
Losses
b
Belief loss
TR.
Local reward loss
Lp
Local transition loss
Combined local loss: C La Li + R LT + +2Kp (Lp Li Lt LO)
C
Lo
Observation loss
L
On-policy reward loss
On-policy transition loss
MDP
a,A
Action, action space
y E [0, 1) Discount factor
M
MDP
P: S A - A(S) Transition function
r, R.R
Reward, reward function
S.S
State, state space
POMDP
6,B
Belief, belief space
T: B x AAx St - B Belief update function
7::
(A- 2)* - B Recursive application of the belief update rule along histories
(A 22)* History, history space
Ha
Stationary distribution over histories based on a policy m
0,
Observation, observation space
O:
A
((2) Observation stochastic function
T: B
A(A) POMDP policy conditioned beliefs
P
POMDP
V- Value function for the policy π
Latent (PO)MDP
:p: B Ω B Belief encoder
* (A. Ω)*
Recursive application of the belief encoder along histories
37

---

Published as a conference paper at ICLR 2024
Sub-belief encoder, that M (8, ((M(3),,,0) sub-belicf, so that (3) = B
3
Ky
1-y
b, B
Latent belief, latent belief space
7: B x A x Q - B Latent belief update
=": (A- -B Recursive application of the latent belief update rule along histories
M
Latent MDP
0: A(2)
Latent observation function
#: B - A(A)
POMDP latent policy conditioned latent beliefs
P
Latent POMDP
P: 3 A- A33 Latent probability transition
S.S
Latent state, latent state space
V
Latent value function for the policy
M
MAF; allows to pass from sub-beliefs to latent beliefs b via M((3) = b B
R,
Latent reward, latent reward function
R.
R 20 S [R(s, a)]
WAE-MDP Components
v: S S x Ω Extended state decoder, defined as (5() (s),
S S - S State decoder
6: Sn S State embedding function
O 5 +00 Deterministic observation decoder
ADDITIONAL. REFERENCES
Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-generation hyperparameter optimization framework. CoRR, abs/1907.10902, 2019. URL
http://arxi::
10902
K. B. Athreya and Ney. A new approach to the limit theory of recurrent markov chains. Trans- actions of the American Mathematical Society, 245:493-501, 1978. ISSN 00029947. URL
http:// www. jstor. wttg//taaleeeee...oee.e httg//taabee111998....
Christel Baier and Joost-Pieter Katoen. Principles of model checking. MIT Press, 2008. ISBN 978-0-262-00649--9
Yoshua Bengio, Nicholas Léonard, and Aaron C. Courville. Estimating propagating gradients through stochastic neurons for conditional computation. CoRR, Abs/1308.3432 2013. URL
http:://rxiv.org/aae
3432
Luca de Alfaro, Thomas A. Henzinger, and Rupak Majumdar. Discounting the future in systems theory. In Jos C. M. Baeten, Jan Karel Lenstra, Joachim Parrow, and Gerhard J. Woeginger (eds.), Automata, Languages and Programming. 30th International Colloquium, ICALP 2003, Eindhoven, The Netherlands, June July 4, 2003. Proceedings, volume 2719 Lecture Notes in Computer Science, pp. 1037 Springer, 2003. doi: 10.1007/3-540-45061-1 httpps:// org
URL
Josée Deshamais, Vineet Gupta, Radha Jagadeesan, and Prakash Panangaden. Metrics for labelled markov processes. Theor. Comput. Sci., 318(3):323-3333 -544 2004, doi: 10.10166f.ess22030000011 URLhttps:: ptattps://ooii wrw.
2003. 0...0....9.....
38

---

Published as a conference paper at ICLR 2024
Rick Durrett. Probability: Theory and Examples, 4th Edition. Cambridge University Press, 2010.
ISBN 9780511179398. doi: 10.0177CB0976611779988. URL /ttt::::/do..og/////
1017 0017/380980011771380
Jiri Fajtl, Vasileios Argyriou, Dorothy Monekosso, and Paolo Remagnino. Latent bernoulli autoen- coder. Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13- 18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learnnng Research, pp. 2964- 2974. PMLR, 2020. URL /tt.:tttp/pprrrcce......... mmm.a.smmmmmm...................
html
Robert Givan, Thomas L. Dean, and Matthew Greig, Equivalence notions and model minimization in markov decision processes, Artif. Intell., 147-222) 2003. doi: 10. (0.1016/5000-3702200)) 00376-4. URL https://cooi.. http:///ooooooo././////.. 10. wt..r//...111/101100
3702 (C2) 00376-4
Ishaan Gulrajani, Faruk Ahmed, Martín Arjovsky, Vincent Dumoulin, and Aaron C. Courville. Improved training of wasserstein gans. In Isabelle Guyon, Ulrike von Luxburg, Samy Ben- gio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 777775777, 2017. URL https: htttps::/proocedigg... nww..ng..se.eu.inee ips htcp/paper/2211//aaeeee html
8 htt.:/acccc22211..cc.......6666.
Ruiquan Huang, Yingbin Liang, and Jing Yang. Provably efficient UCB-type algorithms for learning predictive state representations. In The Twelfth International Conference Learning Represen- tations, 2024.
www.nneeeeweneeeeeeeeeeee www.m JIdSsumXXBdi Jd5PXbBBX
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizzation with gumbel-softmax. In 5th International Conference on Learning Representationss ICLR 2017 Conference Track Pro- ceedings, 11 2017. URL http:: arxii.
Kim Guldstrand Larsen and Ame Skou. Bisimulation through probabilistic testing. Inf. Comput., 94 (1):1-28, 1991. doi: 10. httpp/////axxxv.oooobb////////// wwtpo//gggaa////ooo/ 01144 0890- (91))000---
URL httptt:///c.... wt1.2oorrr1...0 /t.g//1..11..... 1016/
Qinghua Liu, Alan Chung, Csaba Szepesvári, and Chi Jin. When is partially observable reinforce- ment learning not scary? In Po-Ling Loh and Maxim Raginsky (eds.). Conference on Learn- ing Theory, 2-5 July 2022, London, UK, volume 178 of Proceedings of Machine Learning Re- search, pp. 3175-5220 PMIR, 2022. URL pettps://prooee liu22f. ht.ll
Chris J. Maddison, Andriy Mnih, and Yee Whye Tch. The concrete distribution: A continuous relax- ation of discrete random variables. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL /ttp:::////peeeeeeeeeeee net fort//foummmaa d 1 JE5L5g
Alfred Müller. Integral probability metrics and their generating classes of functions. Advances in Applied Probability, 29(2):429 443, 1997. ISSN 00018678. URL http: www. istor.
org/ wttssstablllllee 11abb/11112811
D. Revuz. Markov Chains. North-Holland mathematical library. Elsevier Science Publishers B.V., be/
second (revised) edition, 1984. ISBN 9780444864000. URL httl:httt///boos...tt
Erfan SeyedSalehi, Nima Akbarzadeh, Amit Sinha, and Aditya Mahajan. Approximate informa- tion state based convergence analysis of recurrent leaarning. In Sixteenth European Work- shop on Reinforcement Learning, 2023. URL htttps://ppenreview..eee/// co8W7pvSrS
Bharath K. Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf, and Gert RG Lanckriet. On integral probability metrics.pphi-diveree. phi-divergenness and binary classification. arXiv preprint arXiv:0901 2209,2 2009.
39

---

Published as a conference paper at ICLR 2024
Aäron van den Oord, Oriol Vinyals, and Koray Kavukcuoglu. Neural discrete repre- sentation learning. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Infor- mation Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp- 6310-6315, 2017. URL httpp////prrceeciss......e..oooo
7 03347777630000666663392222202003 Abst http:r/pps.ccppppe////////aaaa/a
Robert J. Vanderbei. Uniform continuity is almost lipschitz continuity. Technical Report SOR-91- 11, Statistics and Operations Research Series, Princeton University, 1991.
40


Reviews:
[{"review_id": "njn64kiYcN", "reviewer": "Reviewer_yC5B", "summary": "The paper studies latent space models for POMDPs, where the authors present theoretical guarantees on the quality of the learned state representation, i.e., capturing the dynamics of the environment, and its suitability for policy optimization. A key ingredient of the proposed formulation is the assumed access to the true state of the environment during the training phase, e.g., available through simulators or additional high fidelity sensors. This assumption enables a latent space model consisting of an augmented MDP with an associated POMDP amenable for learning, where the original POMDP is recovered as a variance parameter defining the latent POMDP observations goes to zero. The authors proceed to establish bounds on the difference in value functions upon conditioning the policy on the latent representation, i.e., the value function of the latent space model vs the true value function, in addition to a Lipschitz-like bound on the value difference in terms of the divergence between the embedded belief states. To evaluate the proposed method, the authors consider three types of POMDPs requiring long-term memory, inferring hidden states, and robustness to noise, where the proposed belief updater shows significant gains.", "questions": "**Missing technical discussion:**\n- Assumption(2) requiring an episodic RL process appears (out of nowhere) in S3.2.\n  - It would help to anticipate the need for this assumption earlier in the introduction\n  - It would also help to discuss how restrictive this assumption is, as was done for Assumption(1), and whether it is mainly a technicality for some of the theoretical derivations.\n  - Perhaps related to this additional assumption, the future work section mentioned adaptations under relaxed assumptions, with a citation to (Lambrechts et al. 2022). It would help to elaborate on some of those potential relaxations.\n\n**Presentation:**\n- Abstract:\n  - Assumptions regarding access to the true state during training should be mentioned in the abstract.\n- Section 3:\n  - It would help to introduce a simple diagram showing the relationship between the different (PO)MDPs, perhaps including the symbols for key new constructs. (I see the authors already have illustrations under Fig.6 on P.20. It appears the authors preferred to include Fig.1, so perhaps readers like me could benefit from at least a pointer to Fig.6 somewhere on P.5.)\n  - (Follow up): I see now I missed how \"bar\" symbols refer to the latent MDP. The definition of $\\bar{V}_{\\bar{\\pi}}$ appears right before Thm.3.2. It would help to establish this notation early in S3.1.\n  - Let me also suggest to break up the paragraph before Eq.3.\n  - Value difference bounds:\n      - The big fraction on Eq.6 reappears on that of Eq.7, with only one extra term in the numerator. Since Thm.3.2 applies to the optimal policy, perhaps it would be possible to make the connection more explicit. (Indeed, this seems to appear halfway through the proof on P.29.)\n      - The variance assumption in Thm.3.3 appears to be a technicality, which I worry it hinders the actual message. Looking at the proof, I see it's only needed for the value function to be almost Lipschitz. Following the pointer to Remark-2 (which only appears *after* Thm.3.3, I see that smoothing the observation in the first place was only introduced to enable training (also alluded to earlier in S3.1). I suggest to amend the statement of Thm.3.3 to clarify how the variance will be made to vanish by design, e.g., *as learning converges to the deterministic mapping ensured by ...*. (I see Fig.4 is obatiend *at the late stage of training*).\n      - (Follow up) Based on the above, I don't see why Remark-2 appears where it is right now.\n\n**Nitpicking:**\n- Abstract: is can be used", "limitations": "", "rating": 10, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Proposes the utilization of extra information available during training to guide the learning.\n- Makes an important contribution to learning of POMDPs, grounded in a clear connection to an underlying MDP, with an effective learning strategy that comes with theoretical guarantees.\n- Excellent exposition with detailed derivations and illustrations.", "weaknesses": "Some assumptions need to be highlighted early on with a clear account of where they may break, or how to workaround them in practice; see below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GTlsHtv5pG", "reviewer": "Reviewer_Rz6s", "summary": "The authors introduce a model-based procedure, the Wasserstein Belief Updater (WBU), for learning/optimizing POMDPs. The components are a latent model learner, a belief update learner, and a policy optimizer. The authors make an assumption that the POMDP latent states can be accessed during training time. They provide some theoretical approximation error bounds which bound the model quality and the representation quality resulting from their losses. The authors provide details on how to practically implement their proposal (Section 4), and provide experimental evaluations for their algorithm (Section 5).", "questions": "There are several technical questions which I have:\n1. What does \"we assume that we have access to the latent model learned by the WAE-MDP\" mean?\n2. What does Lemma 3.1 mean?\n3. In paragraph \"Local losses\", how is one expected to sample from $\\tau^\\star(h)$, the true belief function for history $h$.\n4. What does $\\varepsilon$ represent in Theorem 3.3?\n5. It is interesting that Theorem 3.3. only holds on the optimal policy $\\bar{\\pi}^\\star$. What is the source of this, and can this not be extended to work for other policies of interest?\n6. Figure 3: Belief loss. How can the loss be negative?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper presents a few interesting ideas for learning bisimulation-type metrics for POMDPs, which is a reasonable extension from the MDP case due to the additional structure of needing to approximate the belief states. Some of the ideas are certainly interesting and potentially worth further development. \n- The problem of representation learning in POMDPs is difficult and interesting.\n- The theoretical results do provide some amount of theoretical backing, and at a glance the proofs for the approximation error bounds seem non-trivial due to the partially observable problem setting. I have not checked the proofs very carefully. \n- The ideas and theory are supported by experiments which show some improvement over baselines.", "weaknesses": "- The biggest drawback to this work is the latent observability condition. The methodology being presented is fairly depend on this condition, since it attempts to learn the latent transition model, observation model, and belief states. This simplifies the problem setting quite a bit by making the POMDPs more or less a complicated involved MDP. It is not clear how much applicability this method has beyond the settings where this assumption holds. \n- In terms of related work, there is a fair amount of literature on provable guarantees for exploration in POMDPs under latent observability (e.g. Lee et al. ICML 2023 and related/followup work) or under structural conditions for POMDPs (e.g. Liu et al. 2022 and related/followup work) which is not mentioned here. \n- The current theoretical guarantees are similar to the other papers on learning with bisimulation-like metrics, in that the extent of the theoretical results is a bound on the approximation error incurred via the population loss functions. I would argue that this is not an \"end-to-end\" theoretical guarantee that is useful, since there are many questions about how one can minimize such a loss when only given access to samples from the distribution (i.e. how close are the empirical losses and the population losses, are there any double-sampling issues, is one sampling from distributions one does not have access to, etc.) and if we can minimize it on policies that actually matter for downstream performance rather than only on the on-policy distributions (currently it is simply an \"on-policy\" bound which holds only for the distribution under which the data was collected). As an example of the weakness of such a bound, the usual bisimulation metric in tabular MDPs satisfies a similar bound (Theorem 5.1, Ferns et al. 2004), but there is still an exponential lower bound for learning any bisimulation-like abstraction (Appendix B of Modi et al., 2020). Thus, the argument that this paper presents theory that can be used to give guarantees is either incomplete or inaccurate, depending on one's perspective. The authors could perhaps support their theory by experimentally verifying that the (population) losses are indeed minimized and verifying that the quantities on the LHS of their bounds are controlled by their procedure.\n- In terms of writing and presentation, I feel that more effort should be put into streamlining the notations and explanations. For example, even after several re-reads, it is still difficult to keep track of the notation for the different models MDPs/POMDPs being passed between, to type-check things (e.g. the state embedding functions depend on different things at different points), and to grok which ideas are essential/novel and which are simply details.\n- Similarly, there is a fair amount of statements which are \"fake rigorous\" i.e. do not check out formally. Here is a list:\n 1. \"Unlike in MDPs, stationary policies that are based solely on the current observation of P do not induce any probability space on trajectories of M\". This is probably not what the authors meant to say, since any policy (stationary or not) induces a distribution over trajectories.\n 2. \"This procedure guarantees \\bar{M} to be probably approximately bisimlarly close to M as \\lambda \\rightarrow 0\". Is this a formal claim? There is no PAC bound given here nor in the cited prior work on WAE-MDPs.\n 3. Lemma 3.1: \"There is a well defined probability distribution [...] over histories likely to be perceived at the limit by the agent when it executes \\bar{\\pi} in P\". What does this mean? I can try to guess what this mean, but as a formal statement this does not make sense.\n 4. Theorem 3.3: there is a \"for any \\epsilon > 0\" in the statement and in the bound but no other quantities in the bound depend on $\\varepsilon$\n- The experimental results mild improvement over the baselines in some environments, and no improvement over the baselines in other environments. What is the source of this varying degree of improvement? Can more experiments be performed to understand when this specific approach yields benefits?\n\nReferences:\nFerns, N., Panangaden, P., & Precup, D. (2004). Metrics for finite Markov decision processes. UAI 2004.\nModi A, Jiang N, Tewari A, Singh S. Sample complexity of reinforcement learning using linearly combined model ensembles. AISTATS 2020.\nLee, J., Agarwal, A., Dann, C. and Zhang, T., 2023, July. Learning in pomdps is sample-efficient with hindsight observability. ICML 2023. \nLiu, Q., Chung, A., Szepesvári, C., & Jin, C. (2022, June). When is partially observable reinforcement learning not scary?. In Conference on Learning Theory (pp. 5175-5220). PMLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LUtSUDN3TR", "reviewer": "Reviewer_Gzjw", "summary": "This paper proposes an algorithm that learns the belief in POMDP. \n\n## Learning the Environment\nThe algorithm first learns an estimation of the true MDP via the WAE-MDP method, under assumption that the underlying model in POMDP is accessible in training (*Assumption 1*). The WAE-MDP learns a latent MDP model which encodes not only the true state transition of the POMDP but also the observation mechanism, thus it induces all information of the original POMDP, if it is exactly learned (but WAE-MDP learns an estimation instead).\n\n## Theoretical results\nUnder the episodic assumption of the POMDP (*Assumption 2*),\nthe estimation of the difference between values from the underlying (true) MDP and from the latent MDP is upper bounded theoretically by certain linear combination of local losses and Belief losses, which are expected Wasserstein distances or reward differences between the underlying (true) MDP / belief system and latent MDP / belief system. And in the variance zero limit, choosing a different learning history would result in an upper-bounded difference in optimal Value functions.\n\n## Learning the Belief\nOnce the POMDP is known via the latent MDP, one can use the first-order belief evolution formula derived from Bayesian method to calculate the belief from timestep 0 to any timestep. But Bayesian method requires intractable integration over current belief, so the final method is to train another model to learn the belief. \n\nTraining belief follows online optimization of the objective function which is the expected Wasserstein distance between $\\overline{\\tau}$ and $\\phi$. And due to computational difficulty, expected KL divergence is used as objective instead.\n\n## Experiments\nThe paper tested the algorithm in several environments such as Space Invaders, Stateless Cartpole, etc., compare with R-A2C and DVRL.", "questions": "Questions:\n- In the theoretical part, the existence of $\\mathcal{H}_{\\overline{\\pi}}$ is relying on the ergodicity of an episodic MDP (Huang Bojun, 2020), where the state space $\\mathcal{S}$ is assumed to be discrete (finite or countably infinite). I want to know the topology (and geometry) of the history space. \n\n- Is it the collection of all $(a,o)$ sequences such that $(\\mathcal{A}\\cdot \\Omega)^*=Map(\\mathbb{N}, \\mathcal{A}\\times \\Omega)$, or it is the collection of all finite-length tuples that $(\\mathcal{A}\\cdot \\Omega)^*=\\bigsqcup_{n\\in\\mathbb{N}}(\\mathcal{A}\\times \\Omega)^n$? Categorically speaking, is it a direct product (product) or direct sum (coproduct)? See https://en.wikipedia.org/wiki/Product_(category_theory).\n\n- Is there any theoretical result on the continuous spaces such as CartPole? The definition of ergodicity is slightly different there.\n\n\nTypos\n- I could hardly analyse the second sentence of the abstract, .\n- Page 4, first line is hard to read. (But the meaning is clear)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper provides a solid math ground (theoretical analysis of the POMDP dynamics) by\n    - proving the existence of a distribution on training history set.\n    - uses the distribution above to estimate certain bounds of learning effectiveness (Thm 3.2) and learning variance (Thm 3.3).\nThe results guarantees the learning algorithm would converge to a good approximation.\n\n- Proposed learning methods (objectives) are from theoretical results.\n\n- Outperforms R-A2C and DVRL in certain POMDP environments.", "weaknesses": "- Wasserstein Belief Updater finally chooses the objective function to be KL divergence as a Wasserstein proxy. I agree with the difficulty of optimizing Wasserstein (which is quite singular). But probably the authors can thing about the entropy-regularized version of Wasserstein distance? Hope the work after the following paper could help.\n[ Marco Cuturi, Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances ]\n[ Marco Cuturi, et. al., Fast Computation of Wasserstein Barycenters ]\n\n- I did not see what are the requirements of the sets we discuss? I guess they are at least discrete to make certain Markov chain theoretical results hold. But as the paper talked about Cartpole, I guess the theory would meant to be possibly applied to continuous state spaces (real manifold at least).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "haNDmzYhLC", "reviewer": "Reviewer_Zb4B", "summary": "The paper proposes Wasserstein Belief Updater which is based on the Wasserstein Auto-encoded MDP (WAE-MDP) framework to learn a latent model to update beliefs for partially observable problems. Due to partial observability, in addition to WAE-MDP components, the proposed method learns a latent observation function and a latent belief encoder which recursively updates the beliefs. Experiments show that the RL agent using the updated belief outperforms previous RNN based algorithms.", "questions": "- Can you provide more details on MAF training?\n\n- Can you provide more details on how the sub-belief encoder is trained and why BPTT is not needed with some gradient analysis? \n\n- Some possible typos and missing definitions: \n  - Is $\\varphi^*$ defined somewhere?\n  - In equation (5), should the Wassertein distance be given by $D$ or?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- Every step of the proposed approach is explained well, and their connections to prior work are clearly discussed.\n\n- The WBU approach nicely extends WAE-MDP to the partially observable settings, and the paper also provide theoretical guarantees that extends bisimulation bounds in WAE-MDP to POMDPs.\n\n- The architecture of the latent belief encoder seems to be a novel by having the Masked Auto-Regressive Flows as the output layer on top of a sub-belief encoder. The training of the sub-belief is then done without back-propagation through time.\n\n- Experimental results suggest that the beliefs updated by the latent belief encoder are useful information for decision-making with partial observations. The WBU agent achieves strong performance in several POMDP environments, and it especially outperforms existing algorithms by a significant margin in the environment requiring long-term memorization.", "weaknesses": "- Even though the method is called Wasserstein Belief Updater and the theoretical analysis is for the Wasserstein belief loss, the actual algorithm uses a KL loss for the belief updater part. Hence, the guarantees do not necessarily hold for the actual algorithm.\n\n- The algorithm doesn't have any step for updating the MAF. Is there something missing or is the MAF pre-trained by some other way?\n\n- In Section 4, the authors claim that the training of the sub-belief encoder does not require the challenging BPTT as in RNN. But given the recursive structure of the sub-belief encoder, it is not clear why the claim is true. Even when the sub-belief encoder is updated only by the belief loss, the gradient flow might still need to propagate back if the previous sub-belief is also generated by the same sub-belief encoder. Back-propagation might stop if the previous sub-belief is actually generated by other networks, like using (3), but this part is a bit unclear.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)