Title: Offline RL with Discrete Proxy Representations for Generalizability in POMDPs
Venue: NeurIPS 2023
Accepted: yes
Keywords: Offline RL

POMDP
Abstract: 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 pRoxy 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.

Offline RL with Discrete Proxy Representations for Generalizability in POMDPs

Abstract

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 pRoxy 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.

1 Introduction

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 observabilities 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.

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 observabilities. 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 observabilities. 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 pRoxy 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.

2 Related Works

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.

3 Preliminary

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 , A , O , T, S 0 , O, O 0 , R, H, γ ) , where the underlying process is a Markov Decision Process (MDP) ( S , A , T, S 0 , R, H, γ ) . Concretely, S represents the set of states, A denotes the set of actions, T : S × A × S → [0 , 1] is the transition function (dynamics), S 0 : S → [0 , 1] specifies the initial state distribution, R : S ×A×S → R is the reward function, H ∈ N refers to the time horizon, and γ ∈ [0 , 1) is the discount factor. Additionally, in the POMDP setting, O denotes the set of observations, and O : S × O → [0 , 1] represents the observation function. At the initial time step t = 0 , a starting state s 0 ∼ S 0 ( · ) and an initial observation o 0 ∼ O 0 ( ·| s 0 ) are sampled. At any time step t ∈ 0 , . . . , H -1 , the policy takes an action a t ∈ A in the environment, which updates the state according to the dynamics s t +1 ∼ T ( ·| s t , a t ) . The next observation is then sampled as o t +1 ∼ O ( ·| 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 τ 0: t = ( o 0 , a 0 , o 1 , . . . , a t -1 , o t , a t ) , then the memory-based policy, in its most general form, is defined as π ( a t | τ 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 1 t , s 2 t , . . . , 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 lowdimensional (e.g., tabular data). We introduce an M -dimensional binary mask variable m ∈ M = { 0 , 1 } M × H to signify the observability of each factor: for the t -th time step and the i -th factor, m i t = 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 ( ·| s t ) = [ s i t | m i t = 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 i t = 1 .

MDP

𝑜

!

= 𝑠

!

Diverse POMDPs

Only accessible during training

𝑜

!

= 𝑂

"

(𝑠

!

)

𝑜

!

= 𝑂

#

(𝑠

!

)

𝑜

!

= 𝑂

$

(𝑠

!

)

𝑜

!

= 𝑂

%

(𝑠

!

)

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 π ( a t | τ 0: t ) to maximize the expected discounted return under different masked observation functions: max π E s 0 ∼ S 0 ,m ∼M [ E τ [ ∑ H -1 t =0 γ 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 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 π β ( 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 Φ for generating discrete state representations and then convert the original dataset D = { ( s, a, s ′ , r ) } to a new dataset D z = { ( z, a, z ′ , r ) } , where z = Φ( s ) ∈ Z is the concatenation of representations from different state factors. We will introduce the details of Φ in the following subsection. In the second phase, an arbitrary offline RL algorithm is applied to learn an decision head h ( a | z ) from D z . This head could provide an action for a given state representation z . Combined with Φ , the oracle policy π ( a | s ) , which can be executed on the state space, can be written as π ( a | s ) = h ( a | Φ( s ) ) . However,

Phase 2

:

Training decision head

Decision head

ℎ(*|,)

Original dataset

{(!, *, !

(

, /)}

Full observations

s

Converted dataset

{(,, *, ,

(

, /)}

1

5

7

2

3

4

5

2

5

3

7

3

3

2

1

4

8

2

4

3

1

5

6

6

7

8

8

9

Discrete state representations

-

State encoder

- = Φ (0)

Phase 1

:

Training state encoder and converting dataset

Phase 3

:

Training proxy state encoder

Original dataset

{(!, *, !

(

, /)}

Generated partial observations

Training

Encoding

Training

Proxy state encoder

*

Φ( ̃ -

+

|.

6:+),

, /

+

)

Training

Constructing data

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 ˜ Φ to predict the discrete state representation z from the past history of partial observations instead of the full state information: ˜ z t ∼ ˜ Φ( ·| τ 0: t -1 , o t ) . Here, ˜ 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 ˜ Φ( ·| τ 0: t -1 , o t ) , we obtain a proxy policy , which could address different partially observable scenarios, as ˜ π ( a t | τ 0: t -1 , o t ) = h ( a | ˜ Φ( ·| τ 0: t -1 , o t )) ) . 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 , · · · , s M ] (refer to Section 3). For each factor, we define a codebook e i ∈ R K × D , i ∈ { 1 , 2 , · · · , 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 i j ∈ R D , j ∈ { 1 , 2 , · · · , K } . As shown in Figure 3 (a), the factor encoder ϕ i takes the state factor s i as input and outputs an embedding ˆ 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:

Next, we obtain the discrete compositional representation z by concatenating the discrete representations of all state factors:

Here, we summarize the entire process of mapping s to z as z = Φ( s ) . For simplicity, we refer to Φ as the state encoder , which actually comprises the set of all codebooks and factor encoders: Φ = { e i , ϕ i | i ∈ 1 , 2 , · · · , M } . And its architecture is illustrated in Figure 3 (a).

Learning Objective for Φ . The representation learning objective is based on the Vector Quantized AutoEncoder (VQ-VAE) approach [45] and can be expressed as follows:

Here, Φ represents the set of all codebooks and factor encoders, while Σ denotes the state decoder and Σ( z ) is the state reconstruction. The function sg ( · ) performs a stop-gradient operation to block gradients from entering its argument. The hyperparameter β 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 ˆ 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 ˆ z i . The third term is the commitment loss, which only applies to the factor encoders and encourages ˆ z i to remain close to the selected discrete latent vector z i . We set the value of β to 0.25, as in the original VQ-V AE paper [45]. The training paradigm of state encoder is summarized in Appendix.

4.2 Learning Discrete Proxy Representations over Partial Observations

Masked state encoder

Φ

[89:;]

(/

%

)

/

%

5

%

7

%

7

%#&

̃ -

%

Mask token

!

[&'()]

1

5

7

2

4

Trajectory encoder

"($

#

, &

#$!

, '

#$!

)

Prediction head set

Ω( ̃ ,

!

|3

!

)

6

%#&

1

5

7

2

4

2

5

1

5

7

2

3

4

5

-

Factor encoder

4

&

(0

&

)

1

Codebook

!

!

5

2

K

1

Codebook

!

"

Factor encoder

4

6

(0

6

)

1

5

2

K

5

Concat

0

a

b

Decision head

ℎ(*|,)

t

Trajectory of partial observations

Proxy policy

!

"($|&

=:"#$

, (

"

)

Proxy state encoder

*

Φ( ̃

-|.

?:%#&

, /

%

)

Current full observation

State encoder

Φ(0)

Decision head

ℎ(*|,)

Oracle policy

"($|1)

̃ (~ *

+ (- |/

$:!"#

, 1

!

)

( = +(4)

c

d

State encoder

8(1)

Proxy state encoder

*

8(9 :

"

|&

=:"#$

, (

"

)

In this section, we focus on the alignment of partial observations with the learned discrete state representations derived from full state information. To achieve this, we propose a method that employs a recurrent architecture and a learnable mask token, allowing the model to handle partial observability and estimate the underlying state of the environment. By aligning the estimation process with the learned representations mentioned in the previous subsection, our model can effectively generate proxy state representations from incomplete observations, enabling efficient decisionmaking in partially observable settings. Overall, our aim is to obtain an inference model named proxy state encoder , denoted as ˜ Φ(˜ z | τ 0: t -1 , o t ) , which approximates the true posterior distribution p ( z t | τ 0: t -1 , o t ) . Its architecture is illustrated in Figure 3 (b).

4.2.1 Construction of Partial Observable Training Data

Trajectory of full states

t

t

Trajectory masked by

)

,

Trajectory masked by

)

-

t

The first step to training ˜ Φ is to get a good training dataset. However, one challenge is that we cannot directly collect partial observations from the environment in the offline RL setting. We only have the original offline dataset D = { ( s, a, s ′ , r ) } and the converted representaion dataset D z = { ( z, a, z ′ , r ) } that contains full state information. To address this issue, we construct trajectories of partial observations by utilizing the converted dataset as the training data for ˜ Φ . And we introduce partial observability by randomly setting a portion of mask variables m i t to be 0 , simulating real-world partially observable scenarios. Specifically, we adopt two scenarios :

Dynamical Missing Scenario M 1 : For each time step t , we set the mask variable m i t = 0 at the probability of η , and otherwise at the probability of 1 -η . In this case, the missing factors dynamically change over time with a certain missing ratio.

Factor Reduction Scenario M 2 : Given a trajectory of data whose time length is T , we set the mask variables m i 0: T = 0 at the probability of η , and otherwise at the probability of 1 -η . This scenario arises when part of the sensors are removed, and the corresponding components of the observations are consistently missing over time.

We give examples of trajectories masked by both missing scenarios in Figure 4. And we adopted both scenarios in our training paradigm by using them iteratively, and denoted the process generating random masking variables as m 0: t ∼ M η , which is summarized in Appendix.

4.2.2 Learning Discrete Proxy Representations

After generating training data, the next step to training ˜ Φ is to aggregate information in partial observable trajectory. Concretely, given the partial observation o = [ s i | m i = 1] ,

Here, we denote Φ [ mask ] the masked state encoder, which contains the trained state encoder mentioned in the previous subsection and an additional learnable mask token e [ mask ] ∈ R 1 × D . It directly converts the observable state factor into the learned discrete factor representation, and represents the unobservable state factor as the mask token. Next, we use a trajectory encoder ξ to capture the temporal information in the trajectory of partial observations. Concretely, it takes current observation representations and previous actions as input and output a continuous vector ν t = ξ ( x t , a t -1 , ν t -1 ) .

Discrete State Representation Distribution Estimation. Now, we want to predict the categorical distribution of p ( z t | τ 0: t -1 , o t ) . For each factor i , we build a set of prediction head Ω = { ω i | i = 1 , 2 , · · · , M } which takes the trajectory representation ν t as input, and output the categorical probalility over the discrete embeddings of the unobserved factors in the form of softmax output. And we only estimate the unobservable factor.

During the execution, we predict the proxy state representations as follow:

Learning Objective for ˜ Φ . The overall learning objective for the distribution estimation can be summarized as minimizing the Cross-Entropy loss between the predicted probability distribution and the true discrete embedding associated with the input sample:

The training paradigm of proxy state encoder is summarized in Appendix.

5 Experiments

In this section, we present a thorough empirical evaluation of ORDER, highlighting its effectiveness in improving offline RL algorithms for a variety of partially observable scenarios. We integrate ORDER with the IQL algorithm [28] and assess its performance across an array of environments in the D4RL benchmark [12] under different partial observation situations. Our results clearly indicate that ORDER substantially enhances the generalization performance of policies trained on offline datasets in diverse partially observable conditions. To ensure the robustness of our findings, we conduct all experiments with five distinct random seeds, each consisting of ten separate runs. For detailed information on the architectures and hyperparameters, please refer to Appendix. In the following subsections, we first evaluate the generalization performance of our model and other baselines under varying degrees of partial observability. Subsequently, we carry out an empirical study to analyze the significance of our learned discrete proxy representations.

Dynamical Missing Scenario

5.1 Evaluating Generalization Performance under Diverse Partial Observation Functions

In this subsection, our primary objective is to assess the generalization performance of policies, which are trained on offline datasets containing full observations, under various partial observation scenarios. We evaluate ORDER and other baseline algorithms on gym locomotion tasks and maze navigation tasks in the D4RL benchmark [12]. In our experimental setup, we treat each dimension of the observation space as an independent state factor.

Our proposed model employs the IQL algorithm [28] to train the decision head on the representation space. By combining the trained proxy state encoder with the decision head, our model can make decisions under partial observations. We denote it as IQL_ORDER and compare it with the following baselines: i) IQL with the recurrent architecture ( IQL_R ): We integrate IQL with a state-of-the-art method for online RL in POMDPs [40] to create this baseline. This method utilizes a specialized recurrent neural network architecture to tackle partial observability and can be combined with arbitrary actor-critic algorithms. During training, we employ the same mask strategy (see Section 4.2.1) to construct partial observable training data for it. Furthernore, to indicate which state factors are missing, we use 0 to replace with the value of unobserved factors and concatenate the mask variable m t with the observation vector o t as the extended input. ii) IQL with the strategy of filling zero ( IQL_FZ ): This baseline is trained directly using the IQL algorithm with the original offline dataset. During testing, 0 are used to represent unobservable factors. iii) IQL with the strategy of filling adjacent elements ( IQL_FA ): This baseline is also trained directly using the IQL algorithm with the original offline dataset. During testing, if there is a latest observed value of an unobserved factor, it is used to represent the current unobservable factors; otherwise, 0 is used.

Generalization Performance under Dynamical Missing Scenarios. We first evaluate policy performance under dynamically missing scenarios, where certain state factors may intermittently and randomly be absent over time. We conduct this evaluation across various missing ratios to gauge the models' adaptability. An intrinsic characteristic of a policy with strong generalization ability is its ability to maintain performance as the missing ratio increases. The results are illustrated in Figure 5. IQL_ORDER outperforms other baselines in most cases, demonstrating its superior generalization performance over dynamic missing scenarios. On the other hand, the performance of IQL_ORDER slightly decreases with increasing missing ratios in most cases. However, even in extreme missing scenarios (missing ratio greater than 70% ), it still maintains good performance. Additionally, IQL_FZ and IQL_FA exhibit similar performance, with IQL_FA slightly outperforming IQL_FZ. Their performance is worse than IQL_ORDER, indicating that simply padding zeros or the latest values into unobservable values cannot effectively address this task. Filling adjacent elements proves to be a better strategy than filling zeros. Finally, IQL_R performs the worst in all cases under dynamic missing scenarios. However, it is worth noting that IQL_R performs well in settings where observation functions are single and fixed (see Appendix). This suggests that it struggles to develop effective policies under diverse and dynamic partial observation settings. One possible reason is that the diverse partial observation function leads to enhanced non-stationarity of dynamics, making it challenging to directly train policies on these cases and resulting in highly unstable training.

Generalization Performance under Factor Reduction Scenarios. Next, we evaluate policy performance under factor reduction scenarios, where in each test episode, some state factors are randomly removed throughout the entire episode. Notably, we exclude the IQL_FA model in this task since we cannot obtain the latest values of the missing factors.

The results are depicted in of Figure 6. IQL_ORDER outperforms other baselines in most cases. However, when the missing ratio is high, its performance decreases more significantly than in dynamic changing scenarios. This might be attributed to the increased difficulty of factor reduction scenarios compared to dynamic changing scenarios. In this setting, models cannot infer the current missing factors from their latest observed values; instead, they must rely on other observable factors for inference. Both IQL_FZ and IQL_R continue to perform poorly, indicating that padding zeros is not an effective strategy and that IQL_R cannot achieve stable training under such challenging scenarios.

5.2 Empirical Analysis of Discrete Proxy Representations

In this subsection, we aim to underscore the significance of discrete proxy representations in our framework. We conduct experiments on the Hopper-medium-v2 dataset. For these experiments, we set up two comparison modelsContinuous Reps. and Raw States -both of which employ our three-phase paradigm to develop policies. The Continuous Reps. model uses a Variational Autoencoder (VAE) in the first phase to generate continuous representations, and in the third phase, it predicts Gaussian distributions over these continuous representations. On the other hand, the Raw States model applies an identity function as its state encoder in the first phase and predicts Gaussian distributions of the missing factors in the third phase.

a

Discrete Reps.

Continuous Reps.

Raw States

missing_ratio=10%

missing_ratio=30%

missing_ratio=70%

missing_ratio=90%

Original representations (or states)

Proxy representations (or states)

c

b

d

-

One challenge we would like to address is the potential misalignment between the partial observability encountered during testing and the data on which the RL policy is trained. To examine this, we first sample 10 trajectories generated by different models under various missing ratios in the Factor Reduction Scenarios. We then use t-SNE [46] to visualize their generated proxy representations (or states) and the original representations (or states), which can be regarded as the ground truth under full

observability. The proxy representations are intended to prevent policies from encountering unseen observations, which could lead to severe performance degradation in the offline RL setting. Therefore, a suitable proxy representation should align well with the original representation, even when the missing ratio is high. As shown in Figure 7 (a), we observe that discrete proxy representations align well with the original representations when the missing ratio is below 70% . In contrast, continuous proxy representations begin to mismatch with the original representations when the missing ratio exceeds 10% . Meanwhile, the proxy raw states show a mismatch at a missing ratio of 10% . Furthermore, Figure 7 (b-c) illustrates the generalization performance under both missing scenarios. We find a positive correlation between alignment ability and generalization performance: discrete representations outperform continuous representations, which in turn perform better than raw states. Additionally, discrete representations maintain good performance when the missing ratio is below 70% , but decrease when the missing ratio exceeds 70% in the Factor Reduction Scenario, consistent with the results in Figure 7 (a). These findings emphasize the importance of utilizing discrete proxy representations in our framework for maintaining a strong alignment between proxy and original representations under partial observability. This alignment allows the policy to make better decisions and exhibit improved generalization performance in partially observable settings. We also evaluate the performance of models with different numbers of learned discrete representations, as shown in Figure 7 (d). We find that increasing the number of discrete representations can improve model performance, likely because a higher number of discrete codes enhance the model's expressive ability. However, when the number exceeds 40 , the performance improvement is marginal, indicating that a trade-off between the number of discrete codes and model performance can be achieved.

6 Conclusion

Wepresent ORDER, a novel probabilistic offline RL framework that addresses the challenges of partial observability in real-world scenarios. By leveraging discrete proxy state representations, ORDER significantly improves robustness and generalization performance across diverse observation functions. Our three-stage training approach effectively harnesses the power of discrete state representations and proxy discrete representation alignment. Extensive experiments demonstrate the effectiveness of ORDER in various partially observable settings. However, some limitations still exist. ORDER currently focuses on masked partial observation functions, leaving other forms like perturbations and noises unaddressed. Moreover, its applicability to more complex real-world tasks, such as autonomous driving, remains to be explored. Despite these limitations, we hope ORDER can pave the way for the deployment of RL policy against various partial observabilities in the real world.

References

Andrea Baisero and Christopher Amato. Unbiased asymmetric reinforcement learning under partial observability. In AAMAS , 2022.

Andrea Baisero, Brett Daley, and Christopher Amato. Asymmetric DQN for partially observable reinforcement learning. In UAI , 2022.

Bram Bakker. Reinforcement learning with long short-term memory. In NeurIPS , pages 1475-1482, 2001.

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

Aude Billard and Danica Kragic. Trends and challenges in robot manipulation. Science , 364(6446):eaat8414, 2019.

Konstantinos Bousmalis, Alex Irpan, Paul Wohlhart, Yunfei Bai, Matthew Kelcey, Mrinal Kalakrishnan, Laura Downs, Julian Ibarz, Peter Pastor, Kurt Konolige, Sergey Levine, and Vincent Vanhoucke. Using simulation and domain adaptation to improve efficiency of deep robotic grasping. In ICRA , pages 4243-4250, 2018.

Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. In ICLR , 2023.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01345 , 2021.

Xiaoyu Chen, Yao Mark Mu, Ping Luo, Shengbo Li, and Jianyu Chen. Flow-based recurrent belief state learning for POMDPs. In ICML , pages 3444-3468, 2022.

Jonas Degrave, Federico Felici, Jonas Buchli, Michael Neunert, Brendan Tracey, Francesco Carpanese, Timo Ewalds, Roland Hafner, Abbas Abdolmaleki, Diego de Las Casas, et al. Magnetic control of tokamak plasmas through deep reinforcement learning. Nature , 602(7897):414419, 2022.

Angelos Filos, Panagiotis Tigkas, Rowan McAllister, Nicholas Rhinehart, Sergey Levine, and Yarin Gal. Can autonomous vehicles identify, recover from, and adapt to distribution shifts? In ICML , pages 3145-3153, 2020.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In ICML , pages 2052-2062, 2019.

Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aäron van den Oord. Shaping belief states with generative environment models for RL. In NeurIPS , pages 13486-13498, 2019.

Pengjie Gu, Mengchen Zhao, Chen Chen, Dong Li, Jianye Hao, and Bo An. Learning pseudometric-based action representations for offline reinforcement learning. In ICML , pages 7902-7918, 2022.

Pengjie Gu, Mengchen Zhao, Jianye Hao, and Bo An. Online ad hoc teamwork under partial observability. In ICLR , 2022.

Hongyi Guo, Qi Cai, Yufeng Zhang, Zhuoran Yang, and Zhaoran Wang. Provably efficient offline reinforcement learning for partially observable Markov decision processes. In ICML , pages 8016-8038, 2022.

Abhishek Gupta, Coline Devin, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Learning invariant feature spaces to transfer skills with reinforcement learning, 2017.

David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. NeurIPS , 2018.

Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICML , pages 2555-2565, 2019.

Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. In ICLR , 2021.

Dongqi Han, Kenji Doya, and Jun Tani. Variational recurrent models for solving partially observable control tasks. In ICLR , 2020.

Matthew Hausknecht and Peter Stone. Deep recurrent Q-learning for partially observable MDPs. arXiv preprint arXiv:1507.06527 , 2015.

Nicolas Heess, Jonathan J Hunt, Timothy P Lillicrap, and David Silver. Memory-based control with recurrent neural networks. arXiv preprint arXiv:1512.04455 , 2015.

Stephen James, Paul Wohlhart, Mrinal Kalakrishnan, Dmitry Kalashnikov, Alex Irpan, Julian Ibarz, Sergey Levine, Raia Hadsell, and Konstantinos Bousmalis. Sim-to-real via sim-to-sim: Data-efficient robotic grasping via randomized-to-canonical adaptation networks. In CVPR , pages 12619-12629, 2019.

Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In ICML , 2022.

Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence , 101(1-2):99-134, 1998.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit Q-learning. In ICLR , 2022.

Ilya Kostrikov, Jonathan Tompson, Rob Fergus, and Ofir Nachum. Offline reinforcement learning with Fisher divergence critic regularization. In ICML , pages 5774-5783, 2021.

Aviral Kumar, Justin Fu, George Tucker, and Sergey Levine. Stabilizing off-policy Q-Learning via bootstrapping error reduction. In NeurIPS , pages 11761-11771, 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-Learning for offline reinforcement learning. In NeurIPS , pages 1179-1191, 2020.

Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. NeurIPS , pages 741-752, 2020.

Jonathan N. Lee, Alekh Agarwal, Christoph Dann, and Tong Zhang. Learning in pomdps is sample-efficient with hindsight observability. In ICML , 2023.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint 2005.01643 , 2020.

Bogdan Mazoure, Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Improving zero-shot generalization in offline reinforcement learning using generalized similarity functions. In NeurIPS , 2022.

Lingheng Meng, Rob Gorbet, and Dana Kuli´ c. Memory-based deep reinforcement learning for pomdps. In IROS , pages 5619-5626, 2021.

Piotr Mirowski, Razvan Pascanu, Fabio Viola, Hubert Soyer, Andy Ballard, Andrea Banino, Misha Denil, Ross Goroshin, Laurent Sifre, Koray Kavukcuoglu, Dharshan Kumaran, and Raia Hadsell. Learning to navigate in complex environments. In ICLR , 2017.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature , 518(7540):529-533, 2015.

Hai Huu Nguyen, Andrea Baisero, Dian Wang, Christopher Amato, and Robert Platt. Leveraging fully observable policies for learning under partial observability. In CoRL , 2022.

Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free RL can be a strong baseline for many POMDPs. In ICML , pages 16691-16723, 2022.

Martin L Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming . John Wiley & Sons, 2014.

Anton Maximilian Schäfer and Hans Georg Zimmermann. Recurrent neural networks are universal approximators. In ICANN , pages 632-640, 2006.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of Go with deep neural networks and tree search. Nature , 529(7587):484-489, 2016.

Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In IROS , pages 23-30, 2017.

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. NeurIPS , 30, 2017.

Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of machine learning research , 9(11), 2008.

Yuhui Wang and Xiaoyang Tan. Deep recurrent belief propagation network for POMDPs. In AAAI , pages 10236-10244, 2021.

Andrew Warrington, Jonathan W Lavington, Adam Scibior, Mark Schmidt, and Frank Wood. Robust asymmetric learning in pomdps. In ICML , pages 11013-11023, 2021.

Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. NeurIPS , 28, 2015.

Daan Wierstra, Alexander Foerster, Jan Peters, and Juergen Schmidhuber. Solving deep memory POMDPs with recurrent policy gradients. In ICANN , pages 697-706, 2007.

Kwangjin Yang, Seng Keat Gan, and Salah Sukkarieh. A gaussian process-based RRT planner for the exploration of an unknown and cluttered environment with a UAV. Advanced Robotics , 27(6):431-443, 2013.

Wenshuai Zhao, Jorge Peña Queralta, and Tomi Westerlund. Sim-to-real transfer in deep reinforcement learning for robotics: a survey. In 2020 IEEE Symposium Series on Computational Intelligence, SSCI 2020, Canberra, Australia, December 1-4, 2020 , pages 737-744. IEEE, 2020.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes ORDER, an offline RL solution where the agent have access to states in the offline data, but can only witness observations where some groups of dimensions are masked when deployed due to occlusion or perturbation in real-life scenarios. To address this problem, a 3-step solution is proposed. First, a state encoder is trained to convert the states into discrete representations, which is essentially multiple VQ-VAEs grouped by observability. Then, the RL agent is trained with IQL on the offline dataset with state dimensions randomly blocked, using representations produced by the encoder as observations. Finally, a proxy state encoder is trained to estimate the current discrete representation based on the past trajectory and current state in deploy time, and its output is used for the trained RL agent. On several mujoco testbeds, ORDER is proved to work better than the baselines and gives representation with better alignment to ground truth.\n\nSTRENGTHS:\n1.\tThis paper is clearly written and easy to follow. Figure 1, 2 and 3 effectively summarizes the main idea of the algorithm, helping the reader to understand. The math symbols are organized and the motivation for each component such as the use of discrete representations are well-explained.\n\n2.\tThe solution proposed by this paper is sound, novel and makes much sense, especially the discrete representation part. From my understanding, there could be two advantages of using discrete representation over continuous: 1) the training of VQ-VAE automatically encourages different values of states to be clustered into one representation, and thus becoming more robust to estimation by the proxy representation, and 2) the RL agent works better on the discrete state space.  \n\n3.\tThe experimental results is neat and clearly shows the superiority of ORDER to other baselines; for example, the visualization result in Fig. 6 effectively shows why discrete representation is used.\n\nWEAKNESSES:\n**1.\tThe related work section could be improved.** \n\nA)\tThe reference list at a single point is too long (line 53, 55 and 78), while the occurrence of reference is too few in the section. The referenced related work is not (only) a proof that authors have read many papers, but (also) should summarize the prior works with a taxonomy that helps the reader to understand the main lines of work and the position of this work in the context. For example, in line 78, “both empirical and theoretical studies …” should be “both empirical […] and theoretical […] studies …”; in line 93-97, point i) and ii) should each have a reference list; in line 103-104, a list of example should be made about “specific partial observation” with reference list for each kind.     \n\nB)\tThere should be a brief discussion on sim2real in the robotics community, since the motivation of the paper is to train the agent in a controlled laboratory environment (quoted from line 34) and delopyed in real-world disaster scenarios (quoted from line 37).\n\n**2.\tThe experiment section could be stronger.**\n\nA)\tMore clarification could be made about the second point in the strength. Specifically, is the performance better because of better representations, or better because IQL works better in the discrete state space? \n\nB)\tMore ablation could be added. For example, what if multiple dimensions are in one group of visibility (I assume the hyperparameters for VQ-VAE would change)? What is the effect of the dimension and the number of discrete codes of the codebook of VQ-VAE? What if the dataset is not medium, but rather a mixture of expert and random?\n\n**3.\tThe computational resource consumption, negative societal impact and licenses are not specified.**\n\n**4.\tOther minor problems:**\n \nA)\tLine 285, bseline -> baseline;\n\nB)\tLine 60 in the appendix, GRU should be cited;\n\nC)\tI would suggest the authors to give clearer, more intuitive explanation at the beginning of the paper about why discrete rather than continuous representation is used. Currently, the reason for using discrete representation is only stated at the end of introduction, and is only empirical without intuition.\n\nQUESTIONS:\nI have one question below, apart from those in the weakness section:\n\nFrom line 196, we know that one encoder only takes one state factor into account; meanwhile, from line 14-15 in the appendix, we know that each dimension is seen as a state factor. Does this mean that the embedding is generated assuming that (groups of) dimensions are independent to each other? If it is, then there is a problem: the agent perceives the same environment, which should not be independent on different sources (with different mask of visibility).\n\nBelow are my suggestions for the author. I will be happy to increase my score if the authors address those problems (see weakness sections for details on point 1-4):\n\n1. Improve the related work section; \n\n2. Present more ablation results;\n\n3. Append the computational resource consumption, negative societal impact and licenses;\n\n4. Fix the other minor problems;\n\n5. I strongly advice the authors to open source upon the acceptance (or next submission) of the paper.\n\nLIMITATIONS:\nThe paper does a good job in discussing the limitations, and I am convinced that despite of those limitations, the work is still valuable to the RL community; however, the paper does not mention any potential negative societal impact of the work. For example, while mujoco is still far from real-life applications, works on automated decision making from lab environment to real life is possible to substitute human work and cause job loss.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe work looks to tackle a subset of the partially observable offline reinforcement learning problem setting, where a dataset of offline experience (of full states and masked states) is given during training, and an agent is tested on masked state features during test time. The authors propose the ORDER training framework, where during training time, an agent first learns discrete state representations using the full states, then uses an RNN to learn (in a supervised manner from the discrete representations from the full states) the discrete state representations over the masked states.\n\nWith this framework in place, the authors show that this approach does well in a masked version of the D4RL offline reinforcement learning benchmark, which they create themselves. They show this form of discretization outperforms baselines algorithm performs better on this benchmark as compared to other baselines. They also show an improvement as a higher percentage of the states are masked out.\n\nSTRENGTHS:\nWithin this specific problem setting, the authors do a good job of showing the benefits of their framework and of discretization. It seems that with this problem set up, discretization and supervised learning over masked states helps quite a bit with performance, and is also quite robust to the percentage of masking based on the plots shown in section 5.\n\nWEAKNESSES:\nThere are a few important issues that this paper seems to have, which I’ll describe in broad strokes in this section. For more specific details on these issues, please see the section-by-section review.\n\nThe first big issue is that it seems that the work has mixed their problem setting with their solution method. The two seem innately tied - even the structure of the work mixes the two together. While the authors start the work with describing the general problem setting of partial observability in offline reinforcement learning, the work continues on to describe a problem setting that is quite far off that mark - where the partial observability only stems from simple bernoulli masks, and full state information is given during the training phase. Besides this, additional assumptions are also given to the agent, such as the agent actually knowing *****which***** state features are masked when (which means that **************************************************************************the agent has to know what’s wrong with the state features at every step)**************************************************************************. All of these assumptions on the problem setting are necessary for ORDER to work. While I’m not saying there’s anything fundamentally wrong with the problem setting they describe, **the issue arises with the fact that the proposed solution method, based on the evidence presented in this work, seems to only work in this extremely specific problem setting that they’ve introduced.**\n\nThe language used and structure throughout the work are also quite misleading, especially in the introduction and abstract where it seems the work has overclaimed in the introduction and underdelivered from section 3 onwards. It feels like before reading section 3, you describe your method as general with the ability to tackle POMDPs in an offline setting. But after reading your preliminary section, you seem to reduce the scope suddenly (masked observation functions, different observation functions, access to states during training). I would recommend adding specificity to the problem setting that you describe throughout your work.\n\nBesides this, I find it somewhat suspicious that IQL_R doesn’t seem to work in any of the environments/datasets, and shows the same flat line throughout all environments/datasets. This calls into question implementation or how hyperparameter tuning was done, which was lacking in the appendix.\n\nLastly, I’m also surprised you didn’t reference hindsight state information, seeing as their problem setup is similar to yours. It seems that this setup is more similar to what the authors are trying to tackle.\n\nQUESTIONS:\nHere is a central question I would like answered with regards to the work: From Appendix B, “Following this, we employ a grid search strategy to finalize our choice of hyper-parameters, the details of which are reported in Table 5.” What was this grid search strategy? How did you arrive at these hyperparameters?\n\nBesides this question, here is a section-by-section review/questions with regard to the work:\n\n****1.****\n\n“While POMDP methods excel in online RL, they rely on continuous environment interactions for policy adaptation, making them unsuitable for offline RL” - This statement confused me a bit. Do you mean that you need on-policy samples for POMDP methods?\n\n“we aim to tackle a more general POMDP problem” - This is false. Your problem set up is not more general, seeing as you have access to states in training. The most general class of POMDP problem settings are settings where you don’t have access to state at all. This setting encompasses “POMDP families” by simply expanding the state space instead of saying that there are multiple observation functions.\n\nThis formulation and claim that the authors are making about tackling a “more general POMDP problem” is also **untrue** in your problem setting. There aren’t different observation functions between train and test time. You simply use a stochastic observation function.\n\n****3.****\n\n“In particular, during the training period, the policy has access to the underlying MDP” - I’m assuming you mean *******samples which include state information*******? If you have the underlying MDP then why not do online RL? Please be more precise.\n\n****4.****\n\n“By converting unseen partial observations into discrete forms akin to those present in the training data”\n\n**Is all these machine learning machinery necessary to discretize states?** We’ve had tile coding and different, very effective, state discretization schemes for a while now. Why do we need a VQ-VAE?\n\n“And we introduce partial observability by randomly setting a portion of mask variables m^i_t to be 0, simulating real-world partially observable scenarios.” - this entire problem set up is a big red flag. **The problem setup is not a realistic representation of partial observability ******at all********. From reading this paper, it sounds like you’ve essentially worked backwards, where you’ve defined a constrained solution method first, then defined your problem setting based on this solution method. It seems like your solution method essentially hinges on the fact that your partial observations are of the form of *************independently************* masking your state feature vectors, which is almost never the case in the real world.\n\nDo I understand this correctly that ************************************************************************************************************************the agent is given the masking vector M?********************************************************************************* If this is the case, then this is another huge assumption. The LSTM here is essentially just predicting missing state features, knowing which are missing.\n\n**4.2.2**\n\nWhat is g here?\n\n****5.****\n\n“Our results clearly indicate that ORDER substantially enhances the generalization performance of policies trained on offline datasets in diverse partially observable conditions.” This is claim is tenuous at best. How is this one partial observability setting considered “diverse observable conditions?”\n\nLooking at your appendix, I’m not sure how you’ve swept hyperparameters and arrived at the hyperparameters given for the baselines. Most importantly, the IQL_R baseline. I’m also unsure as to what IQL_R is exactly. Is it simply IQL with an RNN trained over the masked D4RL datasets?\n\n******5.1******\n\n“An intrinsic characteristic of a policy with strong generalization ability is its ability to maintain performance as the missing ratio increases.” - what does strong generalization performance mean? It seems you mean “generalization in terms of number of masked features”, but more specificity in this language is needed, seeing as “generalization” means a whole deluge of different things in the reinforcement learning context. **This seems to also be an issue throughout this work, seeing the context in which generalization is used.**\n\n“This suggests that it struggles to develop effective policies under diverse and dynamic partial observation settings.” - This is markedly untrue in general - it might be the case where **************************************specifically in your problem setting that you invented************************************** this might be true, but I’m also unconvinced this is the case currently as I’m unsure of how you’ve swept and decided on hyperparameters.\n\n******5.2******\n\nIs using the overlap in t-SNE the best idea for seeing overlap in representations? It’s a dimensionality reduction technique, so the overlap doesn’t necessarily mean much. Why not just use a cosine similarity metric instead? It would be better to have numbers to show the “similarities” in the representations.\n\n****6.****\n\n“that addresses the challenges of partial observability in real-world scenarios.” - **this is just flat-out untrue. Mujoco with random Bernoulli masks is far from real-world scenarios.**\n\nLIMITATIONS:\nPlease see weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work study the masked partial observability in reinforcement learning and propose a novel method ORDER to address this challenge. ORDER leverages the alignment of discrete state representations and significantly improves robustness and generalization performance across diverse masked partial observability. Experiments demonstrate the effectiveness of ORDER in some settings.\n\nSTRENGTHS:\nThis work leverages the alignment of discrete state representations learned by the fully observable offline data and discrete proxy representations for the partial observable experiences, novelly constructing a proxy policy for the partial observable online interaction, which improves the robustness and generalization for the diverse partial observation.\n\nWEAKNESSES:\nThe author claims that ORDER is suitable for heterogeneous state factors, but only provides the experiment results for simple homogeneous state factors settings.\n\nQUESTIONS:\n1. Whether it is reasonable to convert any unobservable state factor into a single mask token $e^{[mask]}$? Is there any convincing reason or it is just a practical implementation to enable the algorithm execution? $e^{[mask]}$ is learned by minimizing (6) along with the tray encoder and prediction heads, can you explain what it means to learn such a mask token?\n2. You mentioned that the state factors can be heterogeneous, but it seems that your experiments only include homogeneous cases where each dimension of the observation vector is treated as an independent state factor. I can accept it as a simplified experiment setting, but I wonder that in general heterogeneous cases where each state factor differs in type and size, how to choose the numbers of discrete codes for each factor? A common discretization degree seems unreasonable now.\n3. Do some kind of heterogeneous state factors like high-dimensional multimedia data suitable for discrete representations? Since in your experiments, each state factor is a single scalar, which is an extreme simplification compared to the heterogeneous factors that you claimed, I really doubt this point.\n\nLIMITATIONS:\nORDER currently focuses on masked partial observation, leaving other forms like perturbations and noises unaddressed. Moreover, its applicability to complex heterogeneous high-dimensional real-world tasks remains to be explored.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors present a three stage method for offline RL in POMDPs. It is assumed that the agent has access to full observations (state) during training, but only masked versions of this state (partial observations) during inference or deployment. In the first stage of training, a mapping is learnt from the original state space to a discrete state space. In the second stage, an optimal policy is learnt using offline RL which maps this discrete state to an action. Finally, in the third stage, a proxy state representation model is learnt, which maps the history of partial observations to a proxy discrete state, which is aligned with the discrete state corresponding to the underlying system state. As different state factors can be masked during inference, this method presents a solution that generalises well to a family of POMDPs obtained from the underlying MDP using different masking strategies. The authors demonstrate this method on four MuJoCo environments.\n\nSTRENGTHS:\n1. The paper is novel, very well written and easy to comprehend.\n2. The paper presents generalising an RL solution of an MDP to a family of POMDPs, thereby making RL generalisable and fairly robust to several real-world problems with a type of partial observations (masked state factors).\n\nWEAKNESSES:\n1. There is no analysis/theory presented regarding loss of optimality due to the discrete state representation.\n\nQUESTIONS:\n1. [Typo] In Line 164, shouldn't the serial number be \"i\" instead of \"ii\"?\n2. [Typo] In Eq(1) is a term missing after \"where\"?\n3. How can it be proven that the discrete state obtained is indeed a state sufficient for control?\n4. I did not understand the difference and the motivation between the second and the third terms in Eq (3). Can the authors please explain the same?\n5. \\tau is not defined. Also, the subscript corresponding to \\tau is sometimes just \"t-1\" and sometimes \"0:t-1\". Can these be made consistent with a proper definition?\n6. [Typo] In line 230 consider replacing \"has\" with \"have\".\n7. [Typo] Reference missing (??) in Line 66 of Supplementary Material.\n\nLIMITATIONS:\nThe authors have covered the limitations of their method. While they do not mention the societal impact, since the paper deals with a general algorithm, the societal impact is same as in any RL case.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}