SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pengjie Gu  $^{1,\dagger}$ , Xinyu Cai  $^{1,\dagger}$ , Dong Xing $^{2}$ , Xinrun Wang  $^{1,\ddagger}$ , Mengchen Zhao  $^{3,\ddagger}$ , Bo An $^{1}$   
School of Computer Science and Engineering, Nanyang Technological University, Singapore $^{1}$   
College of Computer Science and Technology, Zhejiang University $^{3}$   
Noah's Ark Lab, Huawei $^{3}$
{pengjie.gu, xinyu.cai, xinrun.wang, boan}@ntu.edu.sg, dongxing@zju.edu.cn, zhaomengchen@huawei.com
Offline Reinforcement Learning (RL) has demonstrated promising results in various applications by learning policies from previously collected datasets, reducing the need for online exploration and interactions. However, real-world scenarios usually involve partial observability, which brings crucial challenges of the deployment of offline RL methods: i) the policy trained on data with full observability is not robust against the masked observations during execution, and ii) the information of which parts of observations are masked is usually unknown during training. In order to address these challenges, we present Offline RL with DiscrEte pProxy representations (ORDER), a probabilistic framework which leverages novel state representations to improve the robustness against diverse masked observabilities. Specifically, we propose a discrete representation of the states and use a proxy representation to recover the states from masked partial observable trajectories. The training of ORDER can be compactly described as the following three steps. i) Learning the discrete state representations on data with full observations, ii) Training the decision module based on the discrete representations, and iii) Training the proxy discrete representations on the data with various partial observations, aligning with the discrete representations. We conduct extensive experiments to evaluate ORDER, showcasing its effectiveness in offline RL for diverse partially observable scenarios and highlighting the significance of discrete proxy representations in generalization performance. ORDER is a flexible framework to employ any offline RL algorithms and we hope that ORDER can pave the way for the deployment of RL policy against various partial observabilities in the real world.
The past decade witnesses the remarkable success of Reinforcement Learning (RL) in various domains, such as Atari games [38], Go [43], and even nuclear fusion [10]. Offline RL, a particularly promising approach, focuses on learning effective policies from previously collected datasets, reducing the need for online exploration and interactions [34]. However, the various partial observables in the real world pose challenges when deploying the learned policy in practice. Examples include autonomous cars facing blocked or perturbed observations [11], drones navigating through cluttered environments with occlusions [51], and robotic manipulators handling objects with varying degrees of visibility [5]. As a result, agents trained on offline datasets must adapt their learned policies to cope with dynamic, uncertain, and partially observable environments, ensuring generalizability in real-world applications.
$^{1}$  Equal contribution.
$^{2}\ddagger$  Corresponding authors.
Our research targets the situations where agents, trained on offline datasets with full state information, encounter various partial observations in real-world deployments. For instance, a robot trained in a lab with complete data on obstacles and victims might face limited visibility due to occlusions, sensor noise, or environmental factors like smoke or varying light in real scenarios. This contrast between training and real-world conditions can hinder the system's generalizability and effectiveness.
There are two main challenges for the generalization against partial observables. i) During testing, the partial observability may differ substantially from the data on which the RL policy was trained, leading to significantly reduced performance. ii) The specific forms of partial observability remain unidentified during training, preventing the tailored training of agents for these observables. These challenges can be understood as the generalization issues within POMDPs. In this paper, our focus narrows to the generalization of a collection of POMDPs arising from an inherent MDP via diverse masked observation functions, which we refer to as a POMDP family.
Existing offline RL methods primarily target fully observable environments with comprehensive state data, causing difficulties in handling real-world scenarios involving partial observability and uncertain observation functions [30, 31, 29, 13, 28, 8, 15, 26, 7, 35]. While POMDP methods excel in online RL, they rely on continuous environment interactions for policy adaptation, making them unsuitable for offline RL [14, 47, 9, 21, 36, 20, 22, 32, 40]. Furthermore, most of these methods are developed for specific POMDPs and thus lack the ability to generalize across different and unknown partial observation functions. Although some attempts have been made to address partial observability in offline RL [17], their effectiveness is constrained by restrictive assumptions, such as a known observation function or a focus on theoretical analysis.
To bridge the research gap, we unveil Offline RL with DiscrEte pProxy representation (ORDER), a framework that leverages discrete proxy state representations for diverse observation functions. Drawing an analogy, consider mapping a vast terrain: A detailed map is challenging to recall, but broadly categorizing into zones simplifies navigation. Similarly, the discrete representation in ORDER categorizes states, aiding RL algorithms. Moreover, even when part of the terrain (state) is obscured, a knowledgeable guide (proxy representation) can infer the underlying area. Our contributions are three-fold. Firstly, we propose a discrete representation of states to improve the robustness against the partial observability and a proxy representation to recover the discrete representation of the state from partial observable trajectory. Secondly, we propose a three-stage training of ORDER: i) we first train the discrete representation on data with full observation through Vector Quantized AutoEncoder (VQ-VAE [45]), ii) we train the decision module based on the discrete state, and iii) we train the proxy discrete representation on the data with various partial observations, aligning with the discrete representation. Finally, we present a thorough evaluation of ORDER, showcasing its effectiveness in enhancing offline RL algorithms for various partially observable scenarios. Our experiments highlight the importance of discrete proxy representations in maintaining alignment under partial observability, leading to improved generalization performance. ORDER is a flexible framework to employ any state-of-the-art (SoTA) offline RL algorithms and we hope ORDER can pave the way for the deployment of RL policy against various partial observabilities in the real world.
Offline RL. Offline RL aims to learn policies from pre-collected transition data, eliminating the need for active data collection [34]. In recent years, this area has seen significant growth in research [30, 31, 29, 13, 28, 8]. Both empirical and theoretical studies have identified overestimation error, caused by out-of-distribution states and actions, as a critical challenge in offline RL [13]. As a result, numerous behavioral regularizations have been proposed to keep the learned policy within the bounds of the offline data during RL training [15, 26, 7, 35]. However, the majority of existing methods focus on fully observable environments, limiting their applicability in real-world situations that often involve partial observability and uncertain observation functions [31, 13, 28, 8, 26]. In such cases, the partial observability encountered during execution can greatly differ from the data on which the RL policy is trained, leading to a significant decline in performance. A noteworthy attempt to tackle partial observability in offline RL is the work by [17]. However, this effort mainly centers on theoretical analysis and assumes a known observation function, which restricts its practical application. Our work sets itself apart from this previous effort by assuming access to full observations in the offline dataset, while the observation functions in the testing environment are unknown, potentially diverse,
and subject to change. Our approach is designed to generalize to uncertain observation functions, thus improving generalization performance across a range of partially observable scenarios.
Online RL for POMDPs. Existing work addressing POMDPs generally falls into two main categories: i) belief state methods [27, 14, 47, 9, 21], which aim to deduce the distribution of hidden (belief) states from partial observation histories and then apply standard RL techniques to these inferred states; and ii) memory-based policies that use memory-based architectures, like recurrent neural networks (RNNs), to store observation trajectories and summarize past experiences [42, 3, 50, 24, 23, 37, 36, 49, 19]. However, most of these methods are developed for specific POMDPs and lack the ability to generalize across different and unknown partial observation functions. Furthermore, these online techniques hinge on the agent's capacity to interact with the environment and adjust its policy online, which is unsuitable for the offline RL setting. Another relevant area is asymmetric reinforcement learning, where existing works assume that agents can access privileged information (e.g., full observations) during training, but need to function in testing environments with partial observations [39, 1, 48, 2, 16, 33]. Nonetheless, they also assume a specific partial observation and depend on data collection through environment interaction. In contrast, our method presumes that we can access full observation information from a fixed dataset without interacting with the environment. Additionally, our method does not focus on a single, fixed partial observation function; it is designed to effectively generalize across a variety of masked observation functions. This adaptability allows our approach to address real-world situations characterized by diverse and uncertain observation functions.
Sim-to-Real RL. Sim-to-real RL is a notable direction in related works. This approach leverages simulation training for real-world applications, adapting policies based on transferred knowledge [52]. A prevalent strategy in this realm is domain randomization [44]. Rather than meticulously replicating real-world parameters, the environment is extensively randomized to encapsulate the genuine data distribution, regardless of simulation-reality bias. Additionally, many researchers adopt domain adaptation, which capitalizes on source domain data to enhance model performance in a target domain, typically data-scarce [6, 18, 25]. A key challenge here lies in unifying the often-divergent feature spaces between source and target domains to facilitate knowledge transfer effectively. Our work shares a thematic connection with the sim-to-real paradigm, as both address the transition from controlled conditions to real-world unpredictability. In our framework, agents, trained comprehensively on offline datasets, must adapt to partial real-world observations. While sim-to-real focuses on bridging the gap between simulation and reality, our emphasis is on transitioning from thorough training data to fragmented real-world observations. Both aim for real-world proficiency, but the challenges and remedies they explore differ substantially.
POMDPs. A Partially Observable Markov Decision Process (POMDP) is a mathematical framework used to model decision-making problems in uncertain environments [41]. Formally, a POMDP is defined as a tuple  $(S, \mathcal{A}, \mathcal{O}, T, S_0, O, O_0, R, H, \gamma)$ , where the underlying process is a Markov Decision Process (MDP)  $(S, \mathcal{A}, T, S_0, R, H, \gamma)$ . Concretely,  $\mathcal{S}$  represents the set of states,  $\mathcal{A}$  denotes the set of actions,  $T: S \times \mathcal{A} \times S \to [0, 1]$  is the transition function (dynamics),  $S_0: S \to [0, 1]$  specifies the initial state distribution,  $R: S \times \mathcal{A} \times S \to \mathbb{R}$  is the reward function,  $H \in \mathbb{N}$  refers to the time horizon, and  $\gamma \in [0, 1)$  is the discount factor. Additionally, in the POMDP setting,  $\mathcal{O}$  denotes the set of observations, and  $O: S \times \mathcal{O} \to [0, 1]$  represents the observation function. At the initial time step  $t = 0$ , a starting state  $s_0 \sim S_0(\cdot)$  and an initial observation  $o_0 \sim O_0(\cdot | s_0)$  are sampled. At any time step  $t \in [0, \dots, H-1]$ , the policy takes an action  $a_t \in \mathcal{A}$  in the environment, which updates the state according to the dynamics  $s_{t+1} \sim T(\cdot | s_t, a_t)$ . The next observation is then sampled as  $o_{t+1} \sim O(\cdot | s_{t+1})$ , and the reward is computed as  $r_{t+1} = R(s_t, a_t, s_{t+1})$ . Let the observable trajectory up to time step  $t$  be denoted as  $\tau_{0:t} = (o_0, a_0, o_1, \dots, a_{t-1}, o_t, a_t)$ , then the memory-based policy, in its most general form, is defined as  $\pi(a_t | \tau_{0:t-1}, o_t)$ , conditioned on the entire history.
In this paper, we focus on the masked partially observable setting, where state factors are obscured or "masked" in a dynamic manner, such as when sensors of a robot are occluded or signals are perturbed. In this context, any state can be represented as a collection of random variables  $s_t = [s_t^1, s_t^2, \ldots, s_t^M]$ . This implies that the state consists of various factors collected from multiple sources.
Each factor  $s^i$  can differ in type and size, ranging from high-dimensional (e.g., multimedia data) to low-dimensional (e.g., tabular data). We introduce an  $M$ -dimensional binary mask variable  $m \in \mathcal{M} = \{0,1\}^{M \times H}$  to signify the observability of each factor: for the  $t$ -th time step and the  $i$ -th factor,  $m_t^i = 1$  if it is observed, and 0 otherwise. Consequently, we redefine  $O$  as the masked observation function: Given the mask value  $m$ , the current observation  $o_t$  can be get from  $O(\cdot | s_t) = [s_t^i | m_t^i = 1]$ . In this context, different mask variables result in distinct POMDPs due to variations in their observation functions. And the mask variable  $m_t$  may change with the time step  $t$ . An MDP can be considered as a special case of a POMDP when the observation  $o$  is always identical to the state  $s$ , i.e., for any time step  $t$  and factor  $i$ ,  $m_t^i = 1$ .
Figure 1: In general POMDP problem, the policy is trained on samples which contain state information and expected to perform well across various POMDPs with distinct observation functions.
Generalization over a POMDP Family. Prior works generally aim to address a specific POMDP, where the policy is trained and tested within the same POMDP. In this paper, we aim to tackle a more general POMDP problem as depicted in Figure 1. In particular, during the training period, the policy has access to samples which include state information. During the testing period, the policy is expected to perform well in various POMDPs that share the same underlying MDP but have distinct observation functions. The objective for such a problem in the masked POMDP setting can be regarded as finding an optimal policy  $\pi (a_t|\tau_{0:t})$  to maximize the expected discounted return under different masked observation functions:  $\max_{\pi}\mathbb{E}_{s_0\sim S_0,m\sim \mathcal{M}}[\mathbb{E}_{\tau}[\sum_{t = 0}^{H - 1}\gamma^t r_{t + 1}|s_0]]$ .
Offline RL. Our work focuses on the offline setting, where the agent learns a policy from a static dataset  $\mathcal{D} = (s,a,s',r)$  without collecting new experience data. The dataset is generated by another policy, referred to as the behavioral policy and denoted by  $\pi_{\beta}(a|s)$ . We allow full access to state information during the training phase by relaxing partial observability constraints. However, during execution, the trained policy must handle diverse partial observation functions. Our method separates representation learning and policy learning, enabling the use of state-of-the-art offline RL algorithms to train a decision head that can make decisions based on the learned representation (see Section 4). In our paper, we employ Implicit Q-learning (IQL) [28] (see details of IQL in Appendix.) as the offline RL algorithm for training our decision module.
4 Offline RL with Discrete Proxy Representations
We introduce ORDER, a framework designed to tackle challenges in partially observable settings for offline RL. Our approach uses two main strategies: i) We learn discrete proxy state representations to handle out-of-distribution states and uncertainties in observation functions. The choice of discrete state representation is deliberate, as it ensures that our policy maintains its generalizability even in situations characterized by uncertainty. By converting unseen partial observations into discrete forms akin to those present in the training data, ORDER could adapt to a variety of uncertain, partially observable situations. ii) We use an estimated proxy representation to connect partial observations with learned discrete state representations. This allows our model to understand the state information even with limited observations, improving accuracy and making the optimization process easier. In the following subsections, we first give an overview of the training process, then explain how we learn discrete state representations and discrete proxy representations for partial observations.
Overview. As illustrated in Figure 2, ORDER is a three-phase paradigm for developing a policy capable of addressing the general POMDP problem. In the first phase, we train a state encoder  $\Phi$  for generating discrete state representations and then convert the original dataset  $\mathcal{D} = \{(s,a,s',r)\}$  to a new dataset  $\mathcal{D}_z = \{(z,a,z',r)\}$ , where  $z = \Phi(s) \in \mathcal{Z}$  is the concatenation of representations from different state factors. We will introduce the details of  $\Phi$  in the following subsection. In the second phase, an arbitrary offline RL algorithm is applied to learn an decision head  $h(a|z)$  from  $\mathcal{D}_z$ . This head could provide an action for a given state representation  $z$ . Combined with  $\Phi$ , the oracle policy  $\pi(a|s)$ , which can be executed on the state space, can be written as  $\pi(a|s) = h\big(a|\Phi(s)\big)$ . However,
Figure 2: The three-phase training paradigm of ORDER, encompassing (1) training a state encoder  $\Phi$ , (2) training a decision head  $h(a|z)$ , and (3) training a proxy state encoder  $\tilde{\Phi}$ .
the full state information is only available during training. During execution, the overall policy cannot be directly used, since some state factors will be missed. To address this issue, in the third phase, we train a proxy state encoder  $\tilde{\Phi}$  to predict the discrete state representation  $z$  from the past history of partial observations instead of the full state information:  $\tilde{z}_t \sim \tilde{\Phi}(\cdot|\tau_{0:t-1}, o_t)$ . Here,  $\tilde{z}_t$  is the predicted state representation which is conditioned on the past history of observations. Combining the decision head  $h(a|z)$  with the proxy state encoder  $\tilde{\Phi}(\cdot|\tau_{0:t-1}, o_t)$ , we obtain a proxy policy, which could address different partially observable scenarios, as  $\tilde{\pi}(a_t|\tau_{0:t-1}, o_t) = h\big(a|\tilde{\Phi}(\cdot|\tau_{0:t-1}, o_t)\big)$ . We summarize the oracle policy and proxy policy in Figure 3 (c-d).
4.1 Learning Discrete Representations over States
In this subsection, we present a method for learning discrete state representations, drawing inspiration from the VQ-VAE approach [45]. Recall that in our setting, we consider the state as a collection of  $M$  heterogeneous state factors:  $s = [s^1, s^2, \dots, s^M]$  (refer to Section 3). For each factor, we define a codebook  $e^i \in \mathbb{R}^{K \times D}$ ,  $i \in \{1, 2, \dots, M\}$ , where  $K$  represents the size of the discrete latent space (i.e., a  $K$ -way categorical), and  $D$  denotes the dimensionality of each latent embedding vector  $e_j^i \in \mathbb{R}^D$ ,  $j \in \{1, 2, \dots, K\}$ . As shown in Figure 3 (a), the factor encoder  $\phi^i$  takes the state factor  $s^i$  as input and outputs an embedding  $\hat{e}^i$ . The discrete latent variable  $z^i$  is then calculated by the nearest neighbor look-up  $g$  using the codebook  $e^i$  as follows:
$$
z ^ {i} = g \left(\hat {e} ^ {i}, e ^ {i}\right) = e _ {j} ^ {i} \quad \text {w h e r e} j = \arg \min  _ {j} \| \hat {e} ^ {i} - e _ {j} ^ {i} \| _ {2}. \tag {1}
$$
Next, we obtain the discrete compositional representation  $z$  by concatenating the discrete representations of all state factors:
$$
z = \Phi (s) = \operatorname {C O N C A T} \left(z ^ {1}, z ^ {2}, \dots , z ^ {M}\right). \tag {2}
$$
Here, we summarize the entire process of mapping  $s$  to  $z$  as  $z = \Phi(s)$ . For simplicity, we refer to  $\Phi$  as the state encoder, which actually comprises the set of all codebooks and factor encoders:  $\Phi = \{e^i, \phi^i | i \in 1, 2, \dots, M\}$ . And its architecture is illustrated in Figure 3 (a).
Learning Objective for  $\Phi$ . The representation learning objective is based on the Vector Quantized AutoEncoder (VQ-VAE) approach [45] and can be expressed as follows:
$$
\mathcal {L} (\Phi , \Sigma) = \mathbb {E} _ {s \sim \mathcal {D}} \left(\left| \Sigma (z) - s \right| _ {2} ^ {2} + \frac {1}{M} \sum_ {i} ^ {M} \left| \operatorname {s g} \left(\hat {e} ^ {i}\right) - z ^ {i} \right| _ {2} ^ {2} + \frac {\beta}{M} \sum_ {i} ^ {M} \left| \hat {e} ^ {i} - \operatorname {s g} \left(z ^ {i}\right) \right| _ {2} ^ {2}\right). \tag {3}
$$
Here,  $\Phi$  represents the set of all codebooks and factor encoders, while  $\Sigma$  denotes the state decoder and  $\Sigma(z)$  is the state reconstruction. The function  $\mathrm{sg}(\cdot)$  performs a stop-gradient operation to block gradients from entering its argument. The hyperparameter  $\beta$  controls the degree to which the code can change. The first term is the reconstruction loss, which ensures that the encoder captures sufficient state information. Note that for the nearest lookup operation in Equation (1), there is no actual gradient, so we use the straight-through estimator [4] to copy gradients from the decoder input  $z^i$
to the encoder output  $\hat{z}^i$ . The second term is the codebook loss, which only applies to the discrete latent vector and encourages the chosen  $z^i$  to be close to the encoder output  $\hat{z}^i$ . The third term is the commitment loss, which only applies to the factor encoders and encourages  $\hat{z}^i$  to remain close to the selected discrete latent vector  $z^i$ . We set the value of  $\beta$  to 0.25, as in the original VQ-VAE paper [45]. The training paradigm of state encoder is summarized in Appendix.
4.2 Learning Discrete Proxy Representations over Partial Observations
Figure 3: The architectures of (a) the state encoder  $\Phi(s)$  and (b) the proxy state encoder  $\tilde{\Phi}(\tilde{z}|\tau_{0:t-1},o_t)$ . And the constructions of (c) the oracle policy  $\pi(a|s)$  and (d) the proxy policy  $\tilde{\pi}(a_t|\tau_{0:t-1},a_t)$ .