Offline RL with Observation Histories: Analyzing and Improving Sample Complexity
Authors: Joey Hong
Anca Dragan
Sergey Levine
Venue: ICLR 2024 poster (2024)
Keywords: offline reinforcement learning
POMDPs
representation learning
Primary area: reinforcement learning

Abstract: Offline reinforcement learning (RL) can in principle synthesize more optimal behavior from a dataset consisting only of suboptimal trials. One way that this can happen is by "stitching" together the best parts of otherwise suboptimal trajectories that overlap on similar states, to create new behaviors where each individual state is in-distribution, but the overall returns are higher. However, in many interesting and complex applications, such as autonomous navigation and dialogue systems, the state is partially observed. Even worse, the state representation is unknown or not easy to define. In such cases, policies and value functions are often conditioned on observation histories instead of states. In these cases, it is not clear if the same kind of "stitching" is feasible at the level of observation histories, since two different trajectories would always have different histories, and thus "similar states" that might lead to effective stitching cannot be leveraged.  Theoretically, we show that standard offline RL algorithms conditioned on observation histories suffer from poor sample complexity, in accordance with the above intuition. We then identify sufficient conditions under which offline RL can still be efficient -- intuitively, it needs to learn a compact representation of history comprising only features relevant for action selection. We introduce a bisimulation loss that captures the extent to which this happens, and propose that offline RL can explicitly optimize this loss to aid worst-case sample complexity. Empirically, we show that across a variety of tasks either our proposed loss improves performance, or the value of this loss is already minimized as a consequence of standard offline RL, indicating that it correlates well with good performance.

OFFLINE RL WITH OBSERVATION HISTORIES: ANALYZING AND IMPROVING SAMPLE COMPLEXITY
Joey Hong
Anca Dragan
Sergey Levine
UC Berkeley
(joey hong, anca, , sergey
ABSTRACT
Offline reinforcement learning (RL) can in principle synthesize more optimal behavior from dataset consisting only of suboptimal trials. One way that this can happen is by "stitching" together the best parts of otherwise suboptimal trajectories that overlap on similar states, to create new behaviors where each individual state is in-distribution, but the overall returns are higher. However, in many interesting and complex applications, such as autonomous navigation and dialogue systems, the state is partially observed. Even worse, the state representation is unknown not easy to define. In such cases, policies and value functions are often conditioned on observation histories instead of states. In these cases, it is not clear if the same kind of "stitching" is feasible at the level of observation histories, since two different trajectories would always have different histories, and thus "similar states" that might lead to effective stitching cannot be leveraged. Theoretically, we show that standard offline RL algorithms conditioned on observation histories suffer from poor sample complexity, in accordance with the above intuition. We then identify sufficient conditions under which offline RI. can still be efficient intuitively, it needs to learn a compact representation of history comprising only features relevant for action selection. We introduce a bisimulation loss that captures the extent to which this happens, and propose that offline RL can explicitly optimize this loss to aid worst-case sample complexity. Empirically, we show that across a variety of tasks either our proposed loss improves performance, or the value of this loss is already minimized as a consequence of standard offline RI, indicating that it correlates well with good performance.
INTRODUCTION
Deep reinforcement learning (RL) has achieved impressive performance in games (Mnih et al., 2013; Silver et al., 2017: AlphaStar, 2019), robotic locomotion (Schulman al., 2015; 2017), and control (Todorov et al., 2012; Haarnoja et al., 2018). A key challenge in the widespread adoption of RL algorithms the need for deploying a suboptimal policy in the environment to collect online interactions, which be detrimental in many applications such recommender systems (Afsar et al., 2021), healthcare (Shortreed et al., 2011: 1; Wang et al., 2018), and robotics (Kalashnikov et al., 2018). Offline RL aims to learn effective policies entirely from an offline dataset of previously collected demonstrations (Levine et al., 2020), which makes it a promising approach for applications where exploring online from scratch unsafe or costly. A major reason for the success of offline RL algorithms is their ability to combine components of suboptimal trajectories in the offline dataset using common states, a phenomenon called "trajectory stitching" (Fu al., 2019a; 2020).
Most offline RL methods are formulated in a Markov decision process (MDP) where the state is fully observed (Sutton and Barto, 2018). However, in many real-world tasks, the state is only partially observed, corresponding to a partially observable Markov decision process (POMDP) (Spaan). For example, in autonomous driving, the robot is limited to information measured by sensors, and does not directly perceive the positions of every car on the road, much less the intentions of every driver. As another example, in dialogue systems, the conversational agent can only observe (potentially noisy and redundant) utterances of the other agents, while their underlying preferences and mental state are hidden. In fact, there is often not even a clear representation or parameterization of "state" (c.g.. what is the space of human intentions or preferences?). Therefore, in such applications, policies must

---

instead be conditioned on all observations thus far - the observation history. Naively, this leads to concerns on the efficiency of existing offline RL algorithms. Offline RL is much less likely to utilize suboptimal behaviors if stitching them requires shared observation histories among them, histories are much less likely to repeat in datasets that are not prohibitively large.
In this work, we aim to answer the following question: When and how can improve the sample efficiency of offline RL algorithms when policies are conditioned on entire observation histories? Given that observation histories make naive stitching very inefficient, we study this question from the lens of when and how we can enable history-conditioned offiin RL to efficiently leverage trajectory stitching. Our focus is on a theoretic analysis of this question, though we also provide simple empirical evaluations to confirm our findings. Theoretically, we first show that in the worst case, naive offline RI. using observation histories can lead to very poor sample complexity bounds. We show that prior pessimistic offline RL algorithms with near-optimal sample complexity guarantees in fully observed MDPs (Rashidinejad et al., 2021; Jin et al., 2021a) fail to learn efficiently with observation histories. We also propose remedy to this, by learning compact representation of histories that contains only the relevant information for action selection. When these representations induce bisimulation metric over the POMDP, we prove that offline RL algorithms achieve greatly improved sample complexity. Furthermore, when existing offline RL algorithms fail to learn such representations, we propose a novel modification that explicitly does so, by optimizing an auxiliary bisimulation loss on top of standard offline RL objective. Empirically, we show - in simple navigation and language model tasks- that when naive offline RL algorithms fail, using our proposed loss in conjunction with these algorithms improves performance; furthermore, we also show that in tasks where existing offline RL approaches already succeed, our loss is implicitly being minimized. Our work provides, to our knowledge, the first theoretical treatment of representation learning in partially observed offline RL, and offers a step toward provably efficient RL in such settings.
RELATED WORK
Offline RL. Offline RL (Lange et al., 2012; Levine et al., 2020) has shown promise in range of domains. To handle distribution shift (Fujimoto et al., 2018; Kumar et al., 2019), many modern offline RL algorithms adopt a pessimistic formulation, learning lower-bound estimate of the value function or Q-function (Kumar et al., 2020; Kostrikov et al., 2021; Kidambi et al., 2020; Yu et al., 2020; 2021). When they work properly, offline RL algorithms should benefit from "trajectory stitching," or combining components of suboptimal trajectories in the data to make more optimal ones (Fu et al., 2019a; 2020). From a theoretical perspective, multiple prior works show that pessimistic offline RI. algorithms have near-optimal sample complexity. under assumptions on the affinity between the optimal and behavior policies (Liu et al., 2020; Rashidinejad et al., 2021; Xie et al., 2021; Jin et al., 2021b). Notably, Xie et al. (2021) show that pessimistic offline RL. algorithms can attain the information-theerrticc lower-bound in tabular MDPs, and Jin et al. (2021b) show a similar result for linear MDPs. In our work, we consider offline RL where policies condition on observation histories.
Our work studies offline RL in POMDPs. A number of prior works on RL in POMDPs have proposed designing models, such as RNNs, that can process observation histories (Zhang et al., 2015; Heess et al., 2015). Other methods instead aim to learn a model of the environment, for example via spectral methods (Azizzadenesheli al., 2016) or Bayesian approaches that maintains belief state over the environment parameters (Ross et al., 2011; 1; Katt et al., 2018). However, such approaches can struggle scale to large state and observation spaces. Igl et al. (2018) propose approximately learning the belief state using variational inference, which scales to high-dimensional domains but does not have any theoretical guarantees. To our knowledge, provably efficient offline RL methods for POMDPs are still relatively sparse in the literature. Recently, Jin et al. (2020) propose estimating the parameters of a tabular POMDP efficiently using the induced observable operator model (Jaeger, 2000), under an undercompleteness assumption between the observations and hidden state. Guo et al. (2022) propose and analyze a similar approach for linear POMDPs. However, these approaches share the same weaknesses as prior methods that rely on spectral methods in that they do not scale beyond linear domains. In our work, we analyze practical offline RL algorithms that work on general POMDPs, and show sufficient conditions on how they can be provably efficient, well as propose new algorithm that satisfies these conditions.
POMDPs.
Representation learning in RL. Motivated by our theoretical analysis of the efficiency of naïve history-basee policies, we propose an approach for learning compact representations of observations histories to improve the efficiency of offline RL in POMDPs. Multiple prior works consider state
2

---

abstraction in MDPs, often by learning low-dimensional representations using reconstruction (Hafner et al., 2019; Watter et al., 2015) or a contrastive loss (van den Oord et al., 2018). Specifically, our work builds bisimulation metrics (Ferns et al., 2012; Castro, 2019), which identify equivalence classes over states based on rewards and transition probabilities. Recently, Zhang et al. (2021) propose learning representations that follow bisimulation-dderved state aggregation to improve deep RL algorithms, and Kemertas and Aumentado- Arratttddn- (2021) propose extensions that improve robustness. The main objective of our work is not to propose new representation learning algorithm, but to identify when offline RL with observation histories can achieve efficient sample complexity in POMDPs. To our knowledge, we are the first to provably show efficient offline RL in POMDPs using theoretical guarantees derived from representation learning.
3 PRELIMINARIES
The goal in our problem setting is to learn a policy ा that maximizes the expected cumu- lative reward in partially observable Markov decision process (POMDP), given by tuple M = (S, A, O,T, r, H1), where S is the state space, A is the action space, is the ob- servation space, A1 is the initial state distribution, and H is the horizon. When action a A is executed at state s E S, the next state is generated according to s' ~ T(-ls, a), and the agent receives stochastic reward with mean r(s, a) [0, 1]. Subsequently, the agent receives observation ~ E.16").
Typically, POMDPs are defined with a state space representation; in practice though, these notoriously difficult to define, and so instead we transform POMDPs into MDPs over observation histories-heneee henceforth called Aibervation-tiisory--DDd (Timmer, 2010). At timestep
[H], define the observation history тı as the sequence of observations and actions THe = [01, Q1, ..,..... ON] Then, an setervatioonhsstory--o of observation histories, and A is the action space, P1 is the initial observation distribution, and H the horizon. When action a A is executed at h E H. the agent observes h h < o' via of ~ P(-IT, a). where @ denotes concatenation, and receives reward with mean r(T,a).
is given by tuple M (H, A. P.r. P1, H). where is the space
The Q-function Q (T, for a policy (-|rr) represents the discounted long-term reward attained by executing given observation history 7 and then following policy т thereafter. QT satisfies the Bellman recurrence: Q" a) = ( = r(T, a) + Bas=T()ir..)
[Q (T', a'))
The value function V considers the expectation of the Q-function over the policy (*(h) = Be-x(-Ir) (T, Meanwhile, the Q-function of the optimal policy Q" satisfies: Q (T, = r(T, 11) [max./ (*
"r',aa' a')], and the optimal value function is V*(() max, Q" (T,
Finally, the expected cumulative reward is given by (Tz) = Ees=or (T))). Note that we do not condition the Q-values nor policy on timestep h because it implicit in the length of
In offline RI., we are provided with dataset D {[T), az, ri, (e,rr,,))) J.V b of size |D] = N. We assume that the dataset D generated i.i.d. from a distribution pu(T, a) that specifies the effective behavior policy Ta(a)r) == u(T,a)/ a)/ Σa ((T, a). In our analysis, we will use n((T, a) to denote the num- ber of times (T, a) appears in D. and P(-17, a) and f(T, a) denote the em- pirical dynamics and reward distribu- tions in D. which may be different from P and r due to stochasticity and sampling error. Finally, as prior work (Rashidinejad et al., 2021; Ku- mar et al., 2022), we define the subop- timality of learned policy से as SubOpt(ā) = ED- [J(m*) [(A)]
State
Dataset Action
Learned Action
Figure 1: learning able to learn that though the grey trajectory unsuc- cessful, a prefix Tr the trajectory is still optimal when stitched with the suffix of blue trajectory 7'.
Illustrative example of trajectory stitching. Here, Q-
Trajectory stitching. Much of how offline RL can learn efficiently lies in its capability to combine components of suboptimal trajectories to deduce better ones, which is called "trajectory stitching". We illustrate this in Figure 1, where trajectory 7 through state s,-1 does not end in positive reward, but does share common state 8t with trajectory r' that does. In MDPs, offline RL using value iteration
3

---

Because (ss) is known to be positive from observing r', offline RL can deduce that taking action a at 8t-1 also has positive value, without explicitly observing it in the dataset. This becomes complicated in observation = E. P(s' 8t a 1 (T.). But (T+) is not known to be positive because Tt has not been observed in the data. This means that, naively, offline RL on observation history MDPs does not seem benefit from trajectory stitching, which may negatively effect how efficiently it can learn from data. We formalize this in Section 4 by proving that offline RL. can have poor worst-case sample complexity in POMDPs.
will learn Q-values: Q(s1-1,0(-1) =
history MDP, as offline RL will now learn Q(
Notation.
Let n. A mic quantity, changing with context. For d-dimensional vectors x; y, ax)) denotes its i-th entry, and define ((z, v) = >, x(i)v(i)² C2 ())(3.)
[n. 1}. Denote = polylog(IO), |A], H, N). We let e. be polylogarith-
4 SHOWING INEFFICIENCY OF OFFLINE RL IN OBSERVATION-
In this section, we show that existing offline RL algorithms with state-of-the sample complexity guarantees in standard MDPs have significantly worse guarantees in observation history MDPs. We consider a class of offline RL algorithms that Iearn pessimistic value functions such that the estimated value bower-bounds the true one, i.e., V- < V for policy л. Practical implementations achieve this by subtracting penalty from the reward, either explicitly (Yu al., 2020; Kidambi al., 2020) or implicitly (Kumar et al., 2020; Kostrikov et al., 2021). We only analyze one such algorithm that does the former, though our findings can likely be extended to general pessimistic offline RL methods.
We consider a meta-algorithm called pessimistic value iteration (PEVI), originally introduced by Jin al. (2021a). This algorithm relies on the construction of confidence intervals c: H. x A R that are high- probabbiiyy bounds on the estimation error of P, Then, pessimistic Q-values are obtained by solving the Bellman recurrence: Q(T,a) a) < f(T, a)- - a) + E. P(o']T, (T))), where values rr) The learned policy is then f(-(7) + arg max, E. Q(T ())))) We give a full pscudocode of the algorithm in Algorithm 2 in Appendix A.I.
are (F) + E. Ô(T (()))))
Prior work has shown that in tabular MDPs, instantiations of PEVI achieve state-of-the-an sample complexity (Rashidinejad et al., 2021). We choose one such instantiation, where confidence intervals c(T, a) are derived using Bernstein's 's inequality:
aTT,a) HV(P(-IT,a) a), V(r + 1) (n(T,a)A
Hf(T, a) (M(T,a) + ()(T,a)A1)
H
The same instantiation was considered by Kumar et al. (2022), and shown to achieve sample com- plexity approaching the cetemmmation-hhhoeeeee lower-bound The additional dependence on H is due to log [ = H polylog([O) : [LA).. However, we can show that in an observation history MDP, the same algorithm has much worse sample complexity bounds.
We first characterizes the distribution shift between the offline dataset distribution Ju(T, a) and the distribution induced by x², given by d" (T, a), via concentrability coefficient C*
Definition 4.1 (Concentrability of the data distribution). Define C" to be the smallest finite constant that satisfies d' (r, (),,a)/(TT,, a) < C* VT E H, a E A.
Intuitively, the coefficient O* formalizes how well the data distribution μTT, a) covers the tuples (T, a) visited under the optimal т², where C" = corresponds to data from ' and increases with distribution shift. C* was first introduced by Rashidinejad al. (2021) but for standard MDPs. Using C*, we can derive the following camplecooplexxtt bound for PEVI in observation history MDP:
Theorem 4.1 (Suboptimality of PEVI in Tabular POMDPs). In a tabular POMDP, the policy * found by PEVI satisfies
C HHHH N
CHHHHH2 N
SubOpt(f) S
We defer our proof. which follows from adapting existing analysis from standard MDPs to observation history MDPs, to Appendix A. Note that dependence on [ makes the bound exponential in the horizon because the space of observation histories satisfies [H] > 10 H. This term arises due to
4

---

encountering observation histories that do not appear in the dataset; without additional assumptions on the ability to generalize to unseen histories, any offiie RL algorithm must incur this suboptimality (as it can only take actions randomly given such histories), making the above bound tight.
5 ANALYZING WHEN SAMPLE--FFICIENG CAN BE IMPROVED
In this section, we show how the efficiency of offline RI. algorithms can be improved by learning representations of observation histories, containing features of the history that sufficienlly capture what is necessary for action selection. We then provide one method for learning such representations based on bisimulation metrics that, when used alongside existing offline RL algorithms, is sufficient to greatly improve their sample complexity guarantees in enseervtionnhisotooe MDPs.
Intuitively, consider that observation histories likely contains mostly irrelevant or redundant infor- mation. This means that it is possible to learn summarizations, such that instead of solving the observation history MDP, it is sufficient to solve a summarized MDP where the states are summa- rizations, actions are unchanged, and the dynamics and reward function are parameterized by the summarizations rather than observation histories. We formalize our intuition into the following; Assumption 5.1. There exists set Z where Z [H), and E > O, such that the summarized MDP (Z, A, P.r, P1, H) satisfies: for every T E HH there exists a 2 E Z satisfying (T)
S E.
The implication of Assumption 5.1 is that we can abstract the space of observation histories into a much more compact space of summarizations, containing only features of the history relevant for action selection. If the state space was known, then summarizations could be constructed as beliefs over the true state. In our case, one practical way of creating summarizations by aggregating observation histories using the distances between their learned representations. Note that these representations may be implicitly learned by optimizing the standard offline RL objective, or they can be explicitly learned via auxiliary representation learning objective. We describe one possible objective in the following section, which enjoys strong theoretical guarantees.
5.1 ABSTRACTING OBSERVATION HISTORIES USING BISIMULATION METRICS
Bisimulation metrics offer one avenue for learning abstractions of the observation history (Ferns et al., 2012; Castro, 2019). While they are not the only way of learning useful representations, these metrics offer strong guarantees for improving the efficiency of learning in standard MDPs, and are also empirically shown to work well with popular off-policy RL algorithms (Zhang et al., 2021) 1). In contrast, we leverage learning bisimulation metrics and show that they can similarly improve the theoretical and empirical performance of offline RL algorithms in observation abservation-history MDPs.
Formally, we define the on-policy bisimulation metric for policy т on an observation-history-Moe MDDP as
d= (T,T'') = (r rr Wi (PT( (T), PF (+)))
where we superscript the reward and transition function by æ to indicate taking an expectation over π. To simplify notation, let d d= be shorthand for the π -bsimmulation metric.
Rather than using the true bisimulation metric, Zhang et al. (2021) showed that it can be more practical to learn an approximation of it in the embedding space. Similarly, we propose learning an encoder の: H + Rdd such that distances (T, -') = 11²(7) <(r')ll approximate the distance under the T* bisimulation metric d' (T, Such an encoder can be learned implicitly by minimizing the standard offline RL objective, or explicitly by via an auxilliary MSE objective: Φ = arg min I|dg. d= 11.
Then, the encoder can be used to compact the space of observation histories H. into space of summarizations by introducing an aggregator 4 H + that maps observation histories to summarizations. Specifically, the aggregator will cluster observation histories that are predicted to be similar under our learned bisimulation metric, i.e., ( (Tr') for T, H if do(T, r') This means that we can approximate the current observation history MDP with a summarized MDP (2,A, P. r, Pt, H). Any practical offline RL algorithm can be used solve for the policy on the summarized MDP, and the policy can be easily evaluated on the original POMDP by selecting actions according 7- (())). In the following section. we show that doing so yields greatly improved sampled complexity guarantees in the original POMDP.
5.2 THEORETICAL ANALYSIS
In Section 4, we showed that applying a naive pessimistic offline RL algorithm (PEVI), which has optimal sample complexity in standard MDPs, to observation-hhitory-MDDD incur suboptimality
5

---

that scales very poorly (potentially exponentially) with horizon H. Here, show that applying the same algorithm to a summarized MDP, which aggregates observation histories based on how similar their learned representations are, achieve greatly improved original Aaservation-history--DDD
guarantees in the the representations induce bisimulation metric.
The first result show relates the value functions under the original observation-history-- MDPP and summarized MDP induced via the summarization function 4b
Lemma 5.1. Ler 4b H Z be learned aggregator that clusters observation histories such that
((7) ((7)) (T, T') Then, the induced summarized MDP (z, A, Ptr, H) satisfies
[V*(r) - (()))v(((())) & H (s+1(ā--o) - (.)
Next, show an improved sample complexity bound than Theorem 4.1 in tabular MDP. We consider the same instantiation of PEVI as in Section 4. However, rather than operating on the raw observation history T, we use the summarization function (T) obtained by learning a bisimulation metric over the space of histories H. We can show that operating on the space of summarizations Z instead of the observation histories leads the following greatly improved bound:
Theorem 5.1 (Suboptimality of PEVI augmented with db in Tabular POMDPs). In a tabular POMDP, the policy ßt found by PEVI on the summarized MDP (2, A, P.r,
P.r,P. Pl, H) satisfies
H N
C²I2 N
+ 2H (s+ 1 d' 11
SubOpt (T) s
Again, we defer full proofs to Appendix A. Here, we see that rather than exponential scaling in horizon H, offline RL now enjoys near optimal scaling, particularly if | < |H).
6 PRACTICAL APPROACH TO IMPROVING OFFLINE RL ALGRITHTM
As described in Section 5, the key component that enables sample-efficient offline RL is the existence of an encoder Φ : H - Rd that learns compact representations of observation his- tories. Specifically, we showed that if the distances between representations under the encoder do())'') = [|((T) - <(r') l/ match the π -isiimiullttoon metric, offline RL algorithms that leverage these representations enjoy better efficiency when required to condition on observation histories. Note that the bound in Theorem 4.1 is a worst-case result. In the general case, even naïve offline RL algorithms might still naturally learn encoders & as part of the standard training process that produce useful representations. In Section 5, we show that one way of measuring the effectiveness of the representations is by how well they induce bisimulation metric. Though this is not necessary condition, we do show that is sufficient for sample-efficient learning. Therefore, in this section, we propose a practical way to train an encoder induce bisimulation metric.
Note that in practice, we cannot naïvely fit dg to the r"-bisimulation metric d", be- cause it assumes knowledge of: (1) the true reward function and observation dynam- ics P of the environment, and (2) the opti- mal policy л² To remedy this, we propose a practical algorithm similar to the one pro- posed by Zhang al. (2021), where an encoder & and policy T.. operating on the embeddings, are trained jointly. To resolve (1), we fit a reward and dynamics model Fa. Pa using dataset D. Then, to resolve (2), we use the leamned policy R rather than optimal n², which intuitively should converge to nT"
Algorithm 1 Offline RL with Bisimulation Learning
Require: Offline dataset
1: Initialize encoders Φ, Ø
2: for i 1, 2.2.... do
3: Train encoder ((0)
4: Train dynamics Fa. Ps: J(F,
5: Train policy F.
Update target encoder: / + (1 a.ō a
7:
Return R.
Formally, given the current learned policy 7o with encoder d, we train with the bisimulation loss on top of the regular offline RL objective, using the following loss function:
27
J(Ø) = a'~a(-1=²)
F((
where z = o(7),2' ((4)),'=0((11 0(7') are the representations from a target network. We choose D to be an approximation of the W-saserssnein distance; in discrete observation settings, total variation

---

CoL dette
OOL Zerr State Walldtion
Method
Filtered BC
CQL
CQL + Bisimulation
Mean Reward
0.28
0.19
0.68
In gridworld environment, Filtered BC takes the path towards the unsafe goal, CQL tries take the path towards loss always takes the correct path towards the safe goal.
Figure 2:
goal but often incorrectly (by going down instead right), and CQL with bisimulation
||P(-12, a) P(-12', a') [Ia and in continuous settings, use W2 (2((((2z, P(-|2/, a²)) Gaussian distributions. Then, we use @ to train our dynamics model of the environment via the following:
(F. E, ET [rr - F(((T)..)))) J(P.o) = =E ETarro d (()())))))))
where again 2' = o(r'). Here, d is either log- ligellkkoihodd mean rrean-squared error if observations discrete or continuous, respectively. Finally, we perform policy improvement on  which conditions on representations generated by é. Note that this be any existing offline RI. algorithm. We detail pseudocode for the meta-algorithm in Algorithm 1.
EXPERIMENTS
Our experimental evaluation aims to empirically analyze the relationship between the performance of offline RL in partially observed settings and the bisimulation loss we discussed in Section Our hypothesis is that, if naïve offline RL performs poorly on a given POMDP, then adding the bisimulation loss should improve performance, and if offline RL already does well, then the learned representations should already induce bisimulation metric, and thus low value of this loss. Note that our theory does not state that naïve offline RL will always perform poorly, just that has a poor worst-case bound, though we hypothesize that successful offline RL runs might still minimize loss as a byproduct of successful learning when they work well. We deseribe the main elements of each evaluation in the main paper, and defer implementation details to Appendix B.
would not expect explicit bisimulation loss to always be necessary,
7.1 TABULAR NAVIGATION
We first evaluate our hypothesis a task involving navigation in x 10 tabular environment similar to gridworld (Fu et al., 2019b). Like grid- world, the environment we consider contains a start (blue) and goal (green) state, and walls (grey) and lava (red) placed in between. We consider sparse reward where the agent earns a reward of 1 upon reaching the goal state; however, if the agent reaches a lava state, then its reward is 0 for the rest of the trajectory. The agent is able to move in either of the four directions (or choose to stay still). To introduce stochasticity in the transition dynamics, there is 20% chance that the agent travels in different direction (that is uniformly sampled) than commanded. Finally, the horizon of each episode is H = 50. Unlike conventional gridworld, the location of the goal state in our environment changes depending on what states the agent visits earlier in the trajectory. The specific layout shown on the left. If the agent takes downwards path from the start state, they will trip switch that turns the goal into the state in the lower right surrounded by lava; otherwise, the goal is the state in the lower left. Because the location of the goal state is unknown and depends from the observation history the agent. Because the goal state in the lower left "safe" (i.e. not surrounded by lava), optimal agent should intentionally trip the switch by going right.
past behavior, it must be inferred
We construct a dataset of size [D] = 5, 000 where 50% trajectories come from policy that moves randomly, and 50% from a policy that primarily takes the path towards the "unsafe" goal state in

---

Method
Mean Reward (Base Task)
Mean Reward (Hard Task)
BC
0.05 + 0.02
0.01 + 0.01
Filtered BC
0.41 0.12
0.12 f 0.05
CQL
0.64 0.17
0.43 + 0.08
CQL + Bisimulation
0.71 + 0.14
0.58 + 0.09
IQL IQL Bisimulation
0.63 + 0.15
0.48 + 0.07
0.70 0.15
0.61 + 0.11
MOPO
0.55 + 0.18
0.41 + 0.12
Table 1: Mean and standard deviation of scores achieved on ViZDoom navigation task.
the lower right. We train three algorithms this dataset, all of which an RNN to process the observation histories: (I) filtered behavior cloning (BC) on the 25% of trajectories the data with highest reward, (2) conservative Q-learning (CQL) (Kumar al., 2020), which is strong offline RL baseline, and (3) CQL augmented with our proposed bisimulation loss.
In Figure 2, show the state-action visitations of policies learned under each algorithm. As expected, the policy learned by filtered BC primarily takes the path towards the unsafe goal state. However, an optimal policy should take the path rightwards that turns the goal into the "safe" one. Both offline RL algorithms attempt to learn such a policy. However, the policy learned by naive CQL sometimes fails realize that must take the rightward path from the start state in order to do so, resulting in high proportion of failed trajectories. This is likely due to the policy failing to infer the correct goal state due to improperly discarding relevant information in its observation history (as RNNs are known to "forget" states that occur far in the past). As we hypothesized, adding a bisimulation loss remedied this issue, and the learned policy successfully takes the optimal path towards the "safe" goal state.
7.2 VISUAL NAVIGATION
Next, consider a much more complex task with image observations. We aim to show that proposed approach improves offline RL performance even when the observation space is large. The task we consider involves navigating maze from first- person pixel observations, namely the "My Way Home" sce- nario in the ViZDoom environment (Kempka al., 2016). In the task, the agent starts in a random room (among 8 total rooms) at random orientation, and is tasked search for piece of armor that is in a specific room. At each step, the agent observes 320 x 240 rendering of its first-person view of the maze, which we cropped and resized to be 80 x 80 in our experiments. The agent has three available actions: (turn feft, turn right, and move corward]. The figure on the left shows the layout and one possible observation by the agent. The reward at each state is -0.0001 except at the location of the armor, where it is +1, and the agent has H 2, 100 timesteps to find the armor. The starting location of the agent is unknown and must be inferred from history.
We construct dataset of [D] 5 107 frames, where 50% of trajectories come from policy that moves randomly, and 50% from a policy trained via A2C (Mnih et al., 2016) on roughly 5 x 10 frames. The policy performs better than random, but still only successfully solves the task 60% of the time. However, we posit that both the random and A2C policies will occasionally behave optimally on different subsets of the maze, trajectory stiching will enable the learning of policy that drastically improves upon both of them. We consider the following baselines, all of which use the same CNN and RNN process the observation histories: (1) behavioral cloning (BC) on the full dataset, (2) filtered BC on the 40% of trajectories in the data with highest reward, (3) conservative Q-learning (CQL) (Kumar et al., 2020), (4) CQL augmented with our proposed bisimulation loss, (5) implicit Q-learning (IQL) (Kostrikov et al., 2021). (6) IQL. with bisimulation loss, and (7) offline model-based policy optimization (MOPO) (Yu et al., 2020). CQL. are IQL. are both sfate-of-the-ar offline algorithms; meanwhile, MOPO is model-based offline algorithm where ensemble of models is used to generate synthetic rollouts, whereas in our proposed approach, a model is also learned but solely used in contrastive loss.
In Table 1, we show the cumulative rewards achieved by each algorithm across 100 independent evaluations. In the "base" task, the agent spawns in random location, and in the "hard" task, the agent always spawns in the room farthest from the goal (blue in above figure). We see that offline RI. greatly outperforms imitation learning in each environment, and that adding our bisimulation
8

---

loss noticeably improves performance. We also that the improvement is greater in the "hard" task, likely because trajectories are longer and learning compact representations is more important. Finally, we observe that using learned model in a bisimulation loss is much more effective than in traditional model-based optimization, which are likely more sensitive to model bias.
7.3 NATURAL LANGUAGE GAME
Our final task is challenging benchmark to test the capabilities of offline RL natural language task. In particular, we aim to learn agents that successfully play the popular game Wordle. We adopt the details from this task from Snell et al. (2023), but provide a summary below. Although this is a relatively simple task, we use real transformer-basse initial evaluation of our hypothesis at a scale similar modern deep networks.
language models to address it, providing an
In the game, the agent tries to guess a 5-letter word randomly selected from a vocabulary of 400 words. Here, the state is the word and is completely unknown to the agent, and actions consists of sequence of 5 letter tokens. After each action, the agent observes a sequence of 5 color tokens, each with one of three "colors" for each letter in the guessed word: "black" means the guessed letter is not in the underlying word, "yellow" means the guessed letter is in the word but not in the right location, and "green" means the guessed letter in the right location. We give a reward of -1 for each incorrect guess and reward of O for a correct guess, at which point environment interaction ends. The agent gets a maximum of H = guesses to figure out the word.
We use dataset of Wordle games played by real humans and scraped from tweets, which was originally compiled and processed by Snell et al. (2023). We train four algorithms that use GPT-2 (with randomly initialized parameters) as a backbone transformer that encodes observa- tion histories. The supervised methods predict actions via imitation learning as an additional head from the transformer: (1) fine-tuning (FT) uses the entire data, and (2) filtered FT uses top- 25% of trajectories. The offline RL methods are: (3) Implicit Language Q-learning (ILQL) (Snell et al., 2023), and (4) ILQL with bisimulation loss.
Method
Wordle Score
Fine-tuning
-2.83 + 0.05
Filtered Fine-tuning
-3.02 + 0.06
ILQL
-2.21 + 0.03
ILQL + Bisimulation
2..1 + 0.03
Table
Mean and standard deviation of scores achieved after training on human Wordle dataset.
We report mean and standard deviation scores of all method across 200 independent evaluations in Table 2. We see that ILQL with bisimulation learning outperforms all other consid- ered approaches, but only marginally over base ILQL. We hy- pothesize that the reason why base IIQQ already performs well on the Wordle task is because standard training is already learn- ing useful representations that induce a bisimulation metric. We assess whether this is true by measuring our bisimulation loss for ILQL both with and without explicit minimization of the loss in Figure across 5 random runs of each algorithm. We notice that ILQL already implicitly minimizes the proposed loss during standard training. This is in line with our hypothesis, though is somewhat surprising, as II.QI. has no awareness of yet still reduces the loss during training.
4.0
ILQL ILQU+ bisimm
3.5
Bisimul ation Loss
3.0
2.0
1.5
1.0
10
20
30
40
50
60
Gradient Updates (x62.5k)
Figure 3: Bisimulation training loss.
8 DISCUSSION
In this paper, we study the effectiveness of offline RL algorithms in POMDPs with unknown state spaces, where policies must utilize observation histories. We prove that because offline RL cannot in the worst case benefit from "trajectory stitching" to learn efficiently in POMDPs, it suffers from poor worst-case sample complexity. However, we also identify that offline RL can actually be provably efficient with suitable representations. Such representations discard features irrelevant for action selection. We show a sufficient condition for this when the representations induce a bisimulation metric. In addition, we show how to improve existing offline RL algorithms by adding a bisimulation loss to enforce the learning of such representations. While we show that learning representations that induce bisimulation metric is sufficient to improve the effectiveness of offline RL with observation histories, it is by no means necessary. A direction for future work is deriving a more nuanced characterization of when useful representations are learned just by standard offline RL training.

---

REFERENCES
Mohammad Mehdi Afsar, Trafford Crump, and Behrouz H. Far. Reinforcement learning based recommender systems: survey. CoRR, abs/2101.061 066286, 2021.
Alckh Agarwal, Sham Kakade, and Lin F Yang. Model-based reinforcement learning with a generative model minimax optimal. In Conference on Learning Theory, pages 67-83. PMLR, 2020.
DeepMind AlphaStar. Mastering the real-time strategy game starcraft ii. URL: https://deemmind
reaa:/sssaaaaaatteit----re-a--ee -tt-stttt-ggmmeeee game-starcraft-it 2019.
Kamyar Azizzadenesheli, Alessandro I Lazaric, and Animashree Anandkumar. Reinforcement learning of pomdps using spectral methods. 29th Annual Conference on Learning Theory. PMLR, 2016.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. CoRR, abs/1911 1.0291,, 2019.
Norman Ferns, Prakash Panangaden, and Doina Precup. Metries for finite markov decision processes. CoRR, abs/1207.411 4144, 2012.
J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning. In arXiv, 2020. URL https://aaxiv org/ http:/df/2200000 07219.
Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine. Diagnosing bottlenecks in deep q-learning algorithms. In Proceedings of the 36th International Conference Machine Learning. PMLR, 2019a.
Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine. Diagnosing bottlenecks in deep Q-learning algorithms. arXiv preprint 1902:1902..105555 2019b.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. arXiv preprint 1 02900, 2018.
Hongyi Guo, Qi Cai. Yufeng Zhang, Zhuoran Yang, and Zhaoran Wang. Provably efficient offline reinforcement learning for partially observable Markov decision processes. Proceedings of the 39th International Conference on Machine Learning, 2022.
Haamoja, Zhou, Abbeel, and Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In arXiv, 2018. URL / pdf/1801 01290. pdd..
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. International conference on machine learning. International Conference on Machine Learning, 2019.
Nicolas Heess, Jonathan J. Hunt, Timothy Lillicrap, and David Silver. Memory-based control with recurrent neural networks. CoRR, abs/1512.00455, 2015.
Maximilian Igl, Luisa M. Zintgraf. Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. 2018.
Herbert Jaeger. Observable operator models for discrete stochastic time series. Neural Computation, 2000.
Chi Jin, Sham M. Kakade, Akshay Krishnamurthy, and Qinghua Liu. Sample-efficient reinforcement learning of undercomplete pomdps. In Advances in Neural Information Processing Systems, volume abs/2006. 12484, 2020.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline r!? In International Conference Machine Learning. PMLR, 2021a.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference Machine Learning, pages 5084-5096. PMLR, 2021b.
10

---

Niels Justesen and Sebastian Risi. Automated curriculum learning by rewarding temporally rare events. CoRR, 0tos/1803.0713/1 2018.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning, pages 651-673, 2018.
Sammie Katt, Frans A. Oliehoek, and Christopher Amato. Learning in pomdps with monte carlo tree search. CoRR, abs/1 806.0066.33 156,0,531, 2018.
Mete Kemertas and Tristan CumeetaddoAAmsttoonn Towards robust bisimulation metric learning. In Advances in Neural Information Processing Systems, 2021.
Michat Kempka, Marek Wydmuch, Grzegorz Runc, Jakub Toczek, and Wojeiech Jaskowski. ViZ- Doom: Doom-based Al research platform for visual reinforcement learning. In IEEE Conference on Computational Intelligence and Games, pages 341-348, Santorini, Greece, Sep 2016. IEEE. The Best Paper Award.
doi: /ttp///G...11..8666.00
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. arXiv preprint arXiv:2005.0595/, 2020.
Ilya Kostrikov, Jonathan Tompson, Rob Fergus, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. arXiv preprint arXiv:2 2021.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy
q-learning via bootstrapping error reduction. In Advances Neural Information Processing Systems, pages 11711-1177,, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint arXiv:2006.044: 04779, 2020.
Aviral Kumar, Joey Hong, Anikait Singh, and Sergey Levine. When should we prefer offline reinforcement learning over behavioral cloning?, 2022.
Sascha Lange. Thomas Gabel, and Martin A. Riedmiller. Batch reinforcement learning. Reinforce- ment Learning, volume 12. Springer, 2012.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint (tX.ii:2000.0064433 2020.
Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch reinforce- ment learning without great exploration. arXiv preprint arXiv:2007.08202, 2020.
Andreas Maurer and Massimiliano Pontil. Empirical bernstein bounds and sample variance penaliza- tion. arxiv preprint arxiv:0907.37400 2009.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pages 1928-1937, 2016.
Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline rein- forcement learning and imitation learning: A tale of pessimism. arXiv preprint 2IV::2/03.122.... 2021.
Tongzheng Ren, Jialian Li, Bo Dai, Simon S Du, and Sujay Sanghavi. Nearly horizon- free offline reinforcement learning. arXiy preprint arXiv:2 14077, 2021.
Stéphane Ross, Joelle Pineau, Brahim Chaib-draa, and Pierre Kreitmann. A bayesian approach for learning and planning in partially observable markov decision processes. Journal of Machine Learning Research, 2011.
11

---

John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889-1897, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy 2017.
optimization algorithms. arXiv preprint arXiv
Susan M Shortreed, Eric Laber, Daniel J Lizotte, T Scott Stroup, Joelle Pineau, and Susan A Murphy, Informing sequential clinical lecision-making through reinforcement learning: an empirical study, Machine learning. 84(1-2):109-133 2011.
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676)3354-3333 359, 2017.
Charlie Snell, Ily Kostrikov, Yi Su, Mengjiao Yang, and Sergey Levine. Offline rl for natural language generation with implicit language q learning. In International Conference on Learning Representations (ICLR), 2023.
Matthijs T. J. Spaan. Purtially Observable Markov Decision Processes, pages 387-414. Springer Berlin Heidelherg. ISBN 978-3-644-27645-33 doi: 10.1007/978-3-64222776----
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT Press, second edition, 2018.
Stephan Timmer. Reinforcement Learning with History Lists. PhD thesis, Universitat Osnabruck, 2010.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-basee -baee control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, 2012.
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR, abs//807.0n 003440 2018
L. Wang, Wei Zhang, Xiaofeng He, and H. Zha. Supervised reinforcement learning with recurrent neural network for dynamic treatment recommendation. Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2018.
Manuel Watter, Jost Tobias Springenberg. Joschka Boedecker, and Martin A. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in Neural Information Processing Systems, 2015.
Tengyang Xie, Nan Jiang, Huan Wang, Caiming Xiong, and Yu Bai. Policy finetuning: Bridging sample-efficient offline and online reinforcement learning. In Advances in Neural Information Processing Systems, 2021.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model- offline policy optimization. arXiv preprint arXiv:2005.13233 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. arXiv preprint arXiv:2102.08363 2021.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In International Conference on Learning Representations (ICLR), 2021.
Marvin Zhang, Sergey Levine, Zoe McCarthy, Chelsea Finn, and Pieter Abbeel. Policy learning with continuous memory states for partially observed robotic control. CoRR, abs/1507.01173 2015.
Yufeng Zhang, Qi Cai, Zhuoran Yang, Yongxin Chen, and Zhaoran Wang. Can -iffferreee and g-lcarning learn representation? a mean-field theory. arXiv preprint arXiv:2006.0476), 2020.
12

---

A PROOFS
In this section, we show the full proofs for the lemmas and theorems described in the main paper.
A.1 PROOF OF THEOREM 4.
In this section, we proof the performance guarantee for PEVI given by the pseudocode in Algorithm We follow the proof of Theorem 4.2 in Kumar et al. (2022), but adapt it for our setting of observation- history-MDPs.
Algorithm PEVI
Require: Offline dataset D, confidence level
1: Compute n(T, a) from D, and estimate F(r, a), P() e
2: Initialize Q a) + 0, (r) + 0, a(T,a)
3: for h= 1,2,....... H do
4: Construct c(T, a), V(T,a) H x A based D. f(r,a) c(T,a) En P(o' r,a) v(T@O)) (Hx arg max, (T, x( 7), Vr E H.
5: Set Q(T,
6: Set f(.
7: Set ((r) < ((T,-) ( r), V7 E.
8:
Return #
Recall from Section 4, that in a tabular POMDP, confidence intervals a(T,a), V(r,a) E H A can be constructed as Equation
A.1.1 TECHNICAL LEMMAS
Lemma A.1 (Bernstein's inequality). Let X, (X,) be i.i.d random variables with values in [0. 1]. and let 8 > O. Then we have
(Var(X)log(2)))
log(2/8) $ 5
P
E
X
X
n
Lemma A.2 (Theorem 4, Maurer and Pontil (2009)). Let X, (X;) with n 2 2 be i.i.d random
variables with values in [0, 1). Define X Σm X, and Var(X) Ei=(( X)². Let > Then we have
(Va log(2/0)
7log(2/0)
IP
E B
X
+ 3(n 1)
6..
Lemma A.3 (Lemma 4, Ren al. (2021)). Let A A2 > 0 be constants. Let Ezo + be function such that ((] < H, Vi and f(i) satisfies the recursion
f(1)
1
Then, we have that f(0) 6(A1 A2).
A.1.2 PESSIMISM GUARANTEE
The first thing want to show is that with high probability, the algorithm provides pessimistic value estimates, namely that (7) < V*(r) for all h e [H] and 7 e To do so, we introduce notion of "good" event, which occurs when our empirical estimates of the MDP are not far from the true MDP. We define E be the event where
HV(P(. r, Vi(T + .))
(P( (((17, a, - P(. P(-IT,a) VA (T
(3)
(ni a) A
(n(T,a) A
13

---

holds for all i [m]] and (T, E H A. With abuse of notation, we let VT.TT @ -) be vector of values of histories of the form 7 @ for E O. We also define &g to be the event where
ajr (F(r,)) r(T,)) S (n(T,a))1) (m(T,a)A1)
(4)
holds for all (T,a).
We want to show that the good event & =1 occurs with high probability. The proof mostly follows from Bernstein's inequality in Lemma A.1 Note that because P(- a), Vi are not independent, we cannot straightforwardly apply Bernstein's 's inequality. We instead use the approach of Agarwal et al. (2020) who, for each state 8, partition the range of (T()) within a modified s- absorbing MDP to create independence from P. The following lemma from Agarwal et al. (2020) is result of such analysis:
Lemma A.4 (Lemma 9, Agarwal et al. (2020)). For any E [H], (T, E H. x A such that aT),a) 2 L. and ò > O, we have
HV(P(. 1T,)), (T(TT)))e
(P(11,,)- P(. P(-1T,))) Pa(ro))
88.
MnT,a)
n(T, a)
Using this, we can show that & occurs with high probability: Lemma A.S. (( 2 - A
Proof. For each and (T, a). if (T,, a) s 1, then equation 3 and equation 4 hold trivially, For a(T,a) 2 2, we have from Lemma A.4 that
HV(P( IT,)), ValT +-)) m(T,a)
P(- [T, (T()) |P( P(:
58.
n(T,
Similarly, we can use Lemma A.2 to derive
f(T,a)e + n(Tra) a(Tra)
P r(T,a)) >
Var(F(r,a)) SP a((T,a) -((,a))]> V 2(n(T, a) 1) 2(n(T, a) -1)
1o,
where we use that Var(F(T, a)) S F(T, a) for [0, rewards, and with slight abuse of notation, let capture all constant factors. Taking the union bound over all and (7, a) yields the desired result. O
Now, we can prove that our value estimates are indeed pessimistic.
Lemma A.6 (Pessimism Guarantee). On event E, we have that Vv(T) V² V* for any step E H and state 7 E H.
Proof. We aim to prove the following for any h and 7: Vn-1(7) s VA(7) S V*(T) & (T) We prove the claims one by one.
(T-1(7) < (T)::: This is directly implied by the monotonic update of our algorithm.
(T(77) < (r): We will prove this via induction. We have that this holds for Vo trivially. Assume holds for 1, then we have
4$ (T) 2 E (T(.a)) P(: T.a) 1n-1(7 @ []
2 E.
- (T,,a) + P(. r,a) Ph- + .-)
- Eg (a()(t))) (f(s. - r(T,))) (P(: - ((.1T,a)) Vh-1(T 0.)
(T(T),
14

---

where we use that
- r((,a)) (T, 2 (f(s,a)-r((,.)))
+ (P(: (, P(: 7()17,a) (T +
under event E.
Finally, the claim of (T) & (T) is trivial which completes the proof of our pessimism guarantee.
A.1.3 PERFORMANCE GUARANTEE
Now, we are ready to derive the performance guarantee from Theorem 4.1. We start with the following value difference lemma for pessimistic offline RL:
Lemma A.7 (Theorem 4.2, Jin et al. (2021a)). On event &, at any step h E [H]. we have
J( 2
(T(Ta)))(TTa),
(5)
A= (T,a)
where di, (T) (Th T, x; for Th. (01, (0,,,,,...,, 0%).
Proof. The proof follows straightforwardly from Jin et al. (2021a) for standard MDPs by simply replacing states with observation histories.
Now, we are ready bound the desired quantity SubOpt(f*)) = ED [J(=*) [()]]. We have
ED - J(F))] ED
- (t))
(6)
= Ep
I1E} Σ (T)
11 E E H. n(T, = 03 (m
+ Ep
: 02
ED l[v7 H, ((T,)"(()) 031{8) Ex()))())
:=A
We bound each term individually. The first is bounded as A S PP s 2
H8 s N for
choice of 8 2 N
Bound on Az. For the second term, have
A 5 Σ ED Hn((T, 0}}
< d' (T, n (7)) = 0}]
$ C Σ
40"10
9N
where that (I( d' di((T, nr" (7)) 01, and that max 11 P(1 p) - a
15

---

Bound Ag. bounded by
What remains is bounding the last term, which we know from Lemma A.7 is
As S 2E
L1v7 e H. n(T,)x(T)) > 0} dT(T,
A (7,0)
Recall that Cn (T(TTa)) a) is given by
HV(P(: T, ( + B.)
Hf(T. a)L nlT,a)
Hi n(T, a)
(T
(TT,a)
We bound the summation of each term separately. For the third term have,
H
He a(T,a)
Hi n(Tra)
ED
Σ Σ
dM(T,
a)
(T
A=1
(m.a)
H
Hi
(T,,TT((F))
(T)) xx(T,x+((T)
1m1
H
H
H
(M(T,T(()))
N
(T))
M/=1
C*
CHHHH/2
N
Here we use Jensen's ' inequality and that Σ d'((T, S (²T((Ta))) for any (r, a). For the second term, similarly have
f(T,a))
Ed
Σ
(T.aa)
V
n(T,a) a)
(r.a)
H Σ
Hi
IE
dt(r,a)
Σ Σ d(G.a))(FF.))
n((T,a)
(r.a)
(r.o)
H3
where we use Cauchy-Schwarz. Finally, we consider the first term of by (T, a)
H
v(P(: T, a, 17 ))
Ep
a)
V
H
n(T,a)
(T,a)
(H(T,a) Hr
Ep
Σ Σ 17,) PT-17, h= (7.a)
n(T,a)
(7,a)
C H2222 N
ΣΣ
Me ()P( a,)), V ())
(T.0)
Similar to what was done in Zhang al. (2020): Ren et al. (2021) for finite-horizon MDPs, we can bound this term using variance recursion for finite-horizon obsrrvation-hhssoy-MDDP. Define
f()) Σ Σ (((.,))(((( (0-
(7)
(7.
16

---

Using Lemma 3 of Ren et al. (2021), have the following recursion:
C CH H
CH1H/
f(0)
1(++1++
(
N
N
where
Ch H
CH
(GGa))((( T,)),, Vn ( ))+
(8)
N
(F.O)
Using Lemma A.3, we can bound S(0) =0 C' + +
Using that for constant c,
H2
N A(0) +
N
H
cC H
CHHHH
+ +ol +0)
N
N
N
c
2cC
H
2
N
have that
4cC* H2
Substituting this back into the inequality for + yields,
(V
C (HHH1
CH
Finally, we can bound
C CHHH2
CHIHHPL N
Ag
N
Combining the bounds for the three terms yields the desired result.
A.2 PROOF OF LEMMA 5.1
Recall that the on-policy bisimulation metric for policy on an observation-hisoor-MMPPP1 is given by: -
d= (T,TT')) (r))+ W, (P* 7), PT +'))
We use the following lemma that states that dl satisfies the following:
Lemma A.8 (Theorem 3, Castro (2019)). Given any two observation histories 5. H. in an obserratiion-histor---DD and policy π,
Proof. The proof follows straightforwardly from Castro (2019); Ferns et al. (2012) for standard MDPs by simply replacing states with observation histories.
Furthermore, recall that we have summarized MDP P,r.p., H) where observation histories are clustered using aggregator &b Let us define the reward function and transition probabilities for policy in the summarized-MDD as
(*(()) = E(O(T)) I r"(C)dE(S),
(ERRGT)
P=(((7') ( = [(((T)) / P=(((>") ()E(())
(ED(T)
17

---

where is a measure on H.
We have,
(T) (() / Pn(d' [T) / PT (
1
S (r*(r) (() + 1/- PT(T ()) (T (+) dE(S)
(CEO(T)
{(((T))
H
sup -V
H
Using Lemma A.2 and the dual formulation of the W1 metric yields
1
/
(r* -
{()(T))
(E)(T)
H -
+
sup (V* - (*((())))
H
d' (F.C)E((()) + H
[(O(T)) (CENCT)
sup (T (T)
$ {(((T)) / (EEMCC) (E)(()) + +1125 - d')
H-1
+
-r)
H
2 + ddo d'
+ H sup V (+((()))
H -
Taking the suprenum of the LHS and solving yields the desired result.
A.3 PROOF OF THEBREM 5.1
The proof follows straightforwardly from noting that
- - = Ep, [v*()). Ep, ( + 2H (s+ c++ )
where the inequality follows from applying Lemma 5.1.
Bounding the first term simply follows from using the same proof as in Scction A.1, except where the space of observation histories H. is now replaced by the space of summarizations Z. This yields the desired result.
18

---

B EXPERIMENT DETAILS
In this section, we provide implementation details of each evaluated algorithm in each of our experimental domains.
B.1 GRIDWORLD EXPERIMENTS
Network architecture. We single-layer RNN with hidden dimension 128 encode observa- tion histories, which consist all the previously visited states (encoded as one-hot vectors). The output of the RNN is fed through single-layer MLP of hidden dimension 256, whose output the representation used to generate the next action (as softmax distribution), as well as train on using the bisimulation loss.
Training details. We use the hyperparameters reported in Table 3.
Hyperparameter
Setting
CQL
0.1
Bisimulation
0.05
Discount factor
0.99
Batch size
Number of updates per iteration
200
Number of iterations
100
Optimizer
AdamW
l.carning
3e-5
Table 3: Hyperparameters used during training.
B.2 VIZDOOM EXPERIMENTS
Network architecture. We use the same convolutional architecture as in Justesen and Risi (2018).
Specifically, we use three-layer CNN with filter sizes of [32, 64, 32] and strides [4, 2, 1]. which produces 32 feature maps of size 7 7. Then, the flattened input size is fed to dense layer size of hidden size 512, then into a three-layer RNN with hidden dimension 512 to encode observation histories. Finally, the output of the RNN is fed through single-layer MLP of hidden dimension 512, whose output is the representation used to generate the next action (as a softmax distribution), as well train using the bisimulation loss.
Training details. We use the hyperparameters reported in Table 4.
Hyperparameter
Setting
CQL &
0.05
Bisimulation
0.02
Discount factor
0.99
Batch size
64
Number updates per iteration
Number iterations
200
Optimizer
AdamW W
Learning
7e-4
Table 4: Hyperparameters used during training.
B.3 WORDLE EXPERIMENTS
Network architecture. We use the GPT-2 small transformer architecture, with the weights ini- tialized randomly. One head of the transformer used to generate the next action (as softmax
19

---

distribution the set of 26 characters). All heads two-layer MLPs with hidden dimension twice that of the transformer's embedding dimension.
Training details. We the hyperparameters reported in Table 5, All algorithms were trained single V100 GPU until convergence, which took less than 3 days.
Hyperparameter
Setting
ILQL r
0.8
ILQL a
0.001
Bisimulation
0.02
Discount factor
0.99
Batch size
1024
Target network update a
0.005
Number of updates per iteration 60
Number of iterations
60
Optimizer
Adam W
Learning rate
le-4
Table 5: Hyperparameters used during training.
20


Reviews:
[{"review_id": "ZjR0Cy0o4A", "reviewer": "Reviewer_HexG", "summary": "The paper focuses on demonstrating that \"trajectory stitching,\" a characteristic of offline RL algorithms in standard MDPs, is absent when these algorithms are applied in POMDPs, where observation histories influence the training and action selection of the policy. It begins by effectively motivating the issue with theoretical insights and subsequently proposes a bisimulation metric. This metric simplifies the observation histories, and this metric is readily usable by any offline RL algorithm.", "questions": "1. Currently the results are presented solely in singleton games. Can the findings of this paper be extended to other environments, particularly where training and testing distributions vary?\n2. What criteria were used to determine the percentage of trajectories to be omitted in Filtered BC?\n3. In reference to Table 2, do the reported results encompass the entire human World dataset, or is it a selected subset? Additionally, how effectively does the method execute in the online Wordle environment? (see https://github.com/Sea-Snell/Implicit-Language-Q-Learning#playing-wordle)", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-motivated and is easy-to-follow.\n2. It thoroughly explores the theory necessary to both - highlight the issues with current methods and introduce the proposed solution.\n3. The introduced bisimulation metric is versatile, easily integrating with popular offline RL algorithms.", "weaknesses": "1. The paper, despite its detailed theoretical sections, lacks comprehensive empirical evidence to substantiate its claims. It’s unclear why the metric was exclusively tested and compared with CQL, omitting comparisons with more recent offline RL baselines such as IQL. Evaluation of this metric in a continuous action space environment, compared against IQL or XQL, would be beneficial.\n\n2. The observed performances of CQL+Bisimulation and vanilla CQL in VizDoom are notably similar, as is the case with ILQL with and without bisimulation in Section 7.3. Does this imply that the bisimulation metric primarily enhances performance in less complex environments like GridWorld?", "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": "zLAnPPgTbZ", "reviewer": "Reviewer_Hq7u", "summary": "The authors focus on the sample efficiency issue of  offline RL in the setting of POMDP with observation histories. Based on the observaton that two different trajectories would always have different histories, which makes stitching become much harder because no real state can be observed,  the authors propose to learn a more  realistic belief state using the bismulation loss and show that this may significantly reduce the sample complexity when working in this space because there is no need to consider complexity introduced by the observation sequence any more.", "questions": "1) The reason that the bound of the sub-optimality of the learnt policy listed in Theorem 5.1 is much tighter than the bound listed in Theorem 4.1 is mainly due to the fact that the space sizes follow $|Z|<<|H|$, which is assumed. However, the state represenation space could still be very large. In addition, this bound does not work in non-tabular POMDPs.\n\n2) The loss function $J(\\phi)$ of the proposed method aims to clone the bisimulation metric $|\\hat{s}(z,a)-\\hat{r}(z',a')|-D(\\hat{P}(\\cdot|z,a),\\hat{P}(\\cdot|z',a'))$ . However, the target is approximately estimated via $\\hat{r}$ and the dynamics $\\hat{P}$, which in turn are estimated based on the offline dataset D. Hence in essential, the proposed method belongs to model-based offline RLs and should compare the proposed method to current model based offline RL methods.\n\n3) The comparison study conducted in Table 1 may not be fair, because the three methods listed are not designed to learn from observation histories, but they can only get the information from the current sampled tuples. In this sense, the good  performance of the proposed method may due to its capturing additional information from the historical data through state abstraction. \n\n4) We  suggest the authors  attach more experiments on some widely used standard benchmarks, such as MuJoCo suites, to evaluate the effectiveness and feasibility of the proposed methods.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This paper show the usefulness of state abstraction in the setting of offline RL.", "weaknesses": "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": "ukWWR7ITJH", "reviewer": "Reviewer_dUhS", "summary": "This paper aims to implement offline reinforcement learning through better trajectory stitching. The main challenge here is that the problems being looked at involve sequences of observations rather than clear states. Because these states aren't well-defined, it's not easy to link similar ones together, making the \"stitching\" hard. To address this, the authors present an approach to learn a compact representation of history comprising only features relevant to action selection. They adopt a bisimulation loss to do this and show that with this approach, offline reinforcement learning can be more efficient in terms of the worst-case sample complexity.", "questions": "Can you explain how your solution might work with other offline RL methods that don't focus much on trajectory stitching (for example models with a pessimistic formulation mentioned in the related works)?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is written in a very clear and understandable manner.\n\n2. The concept of stitching together trajectories based on observations rather than states is quite applicable, especially since in many real-world situations, we don't have direct access to the game's state.\n\n3. The proposed solution is strongly backed by thorough theoretical analysis.", "weaknesses": "1. This solution appears to be heavily dependent on the concept of trajectory stitching, which might limit its usefulness in other offline reinforcement learning (RL) scenarios.\n\n2. I’m a bit uncertain about the novelty of training an offline RL model based on observation histories. It seems that in most offline RL research, what an agent can access is already observations (rather than states) from previous data collections.", "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)