SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 decision-making in partially observable settings. Overall, our aim is to obtain an inference model named proxy state encoder, denoted as  $\tilde{\Phi} (\tilde{z} |\tau_{0:t - 1},o_t)$ , which approximates the true posterior distribution  $p(z_{t}|\tau_{0:t - 1},o_{t})$ . Its architecture is illustrated in Figure 3 (b).
4.2.1 Construction of Partial Observable Training Data
Figure 4: Illustration of Both Missing Scenarios. The figure showcases trajectories spanning various time steps (horizontal axis) and state factors (vertical axis). Colorful grids represent observable state factors, while grey grids signify masked state factors. The left panel presents the trajectory of the full states. In contrast, the middle and right panels display trajectories following masking by the dynamical missing scenario  $\mathcal{M}_1$  and the factor reduction scenario  $\mathcal{M}_2$ , respectively.
The first step to training  $\tilde{\Phi}$  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  $\mathcal{D} = \{(s,a,s',r)\}$  and the converted representation dataset  $\mathcal{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  $\tilde{\Phi}$ . And we introduce partial observability by randomly setting a portion of mask variables  $m_t^i$  to be 0, simulating real-world partially observable scenarios. Specifically, we adopt two scenarios:
- Dynamical Missing Scenario  $\mathcal{M}_1$ : For each time step  $t$ , we set the mask variable  $m_t^i = 0$  at the probability of  $\eta$ , and otherwise at the probability of  $1 - \eta$ . In this case, the missing factors dynamically change over time with a certain missing ratio.
- Factor Reduction Scenario  $\mathcal{M}_2$ : Given a trajectory of data whose time length is  $T$ , we set the mask variables  $m_{0:T}^i = 0$  at the probability of  $\eta$ , and otherwise at the probability of  $1 - \eta$ . 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} \sim \mathcal{M}_{\eta}$ , which is summarized in Appendix.
4.2.2 Learning Discrete Proxy Representations
After generating training data, the next step to training  $\tilde{\Phi}$  is to aggregate information in partial observable trajectory. Concretely, given the partial observation  $o = [s^i |m^i = 1]$ ,
$$
x = \Phi^ {[ \text {m a s k} ]} (o) = \operatorname {C O N C A T} \left(x ^ {1}, x ^ {2}, \dots , x ^ {M}\right) \text {w h e r e} x ^ {i} = \left\{ \begin{array}{l l} z ^ {i} = g ^ {i} \left(\phi^ {i} \left(s ^ {i}\right)\right) & \text {i f} m ^ {i} = 1, \\ e ^ {[ \text {m a s k} ]} & \text {i f} m ^ {i} = 0. \end{array} \right. \tag {4}
$$
Here, we denote  $\Phi^{[\mathrm{mask}]}$  the masked state encoder, which contains the trained state encoder mentioned in the previous subsection and an additional learnable mask token  $e^{[\mathrm{mask}]} \in \mathbb{R}^{1 \times 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  $\xi$  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  $\nu_{t} = \xi(x_{t}, a_{t-1}, \nu_{t-1})$ .
Discrete State Representation Distribution Estimation. Now, we want to predict the categorical distribution of  $p(z_{t}|\tau_{0:t - 1},o_{t})$ . For each factor  $i$ , we build a set of prediction head  $\Omega = \{\omega^i |i = 1,2,\dots ,M\}$  which takes the trajectory representation  $\nu_{t}$  as input, and output the categorical probability 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:
$$
\widetilde {z} = \operatorname {C O N C A T} \left(\widetilde {z} ^ {1}, \widetilde {z} ^ {2}, \dots , \widetilde {z} ^ {M}\right) \text {w h e r e} \left\{ \begin{array}{l l} \widetilde {z} ^ {i} = z ^ {i} & \text {i f} m ^ {i} = 1 \\ \widetilde {z} ^ {i} \sim \omega^ {i} (\cdot | \nu_ {t}) & \text {i f} m ^ {i} = 0 \end{array} \right. \tag {5}
$$
Learning Objective for  $\tilde{\Phi}$ . 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:
$$
\mathcal {L} \left(e ^ {\left[ \operatorname {m a s k} \right]}, \xi , \Omega\right) = - \underset {t \sim U (0, H - 1), z _ {0: t}, a _ {0: t - 1} \sim \mathcal {D} _ {z}, m _ {0: t} \sim \mathcal {M} _ {\eta}} {\mathbb {E}} \left[ \frac {1}{M} \sum_ {i} ^ {M} \left(1 - m _ {t} ^ {i}\right) \log \omega^ {i} \left(\tilde {z} ^ {i} = z ^ {i} | \nu_ {t}\right) \right]. \tag {6}
$$
The training paradigm of proxy state encoder is summarized in Appendix.
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.
Figure 5: Generalization performance of different models under dynamic missing scenarios.
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. Furthermore, 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.
Figure 6: Generalization performance of different models under factor reduction scenarios.
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 models—Continuous 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.
- Original representations (or states)  $\bullet$  Proxy representations (or states)
Figure 7: Empirical Analysis of discrete proxy representations. (a) T-SNE visualization showing alignment of discrete, continuous, and raw state representations with original representations under varying missing ratios in Factor Reduction Scenarios. (b-c) Generalization performance under dynamic missing and factor reduction scenarios. (d) Performance evaluation with varying numbers of learned discrete representations.
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 enhances 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.
We present 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.
This research is supported by the National Research Foundation, Singapore under its Industry Alignment Fund – Pre-positioning (IAF-PP) Funding Initiative. Any opinions, findings and conclusions or recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore.
[1] Andrea Baisero and Christopher Amato. Unbiased asymmetric reinforcement learning under partial observability. In AAMAS, 2022.
[2] Andrea Baisero, Brett Daley, and Christopher Amato. Asymmetric DQN for partially observable reinforcement learning. In UAI, 2022.
[3] Bram Bakker. Reinforcement learning with long short-term memory. In NeurIPS, pages 1475-1482, 2001.
[4] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
[5] Aude Billard and Danica Kragic. Trends and challenges in robot manipulation. Science, 364(6446):eaat8414, 2019.
[44] 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.
[45] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. NeurIPS, 30, 2017.
[46] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-SNE. Journal of machine learning research, 9(11), 2008.
[47] Yuhui Wang and Xiaoyang Tan. Deep recurrent belief propagation network for POMDPs. In AAAI, pages 10236-10244, 2021.
[48] Andrew Warrington, Jonathan W Lavington, Adam Scibior, Mark Schmidt, and Frank Wood. Robust asymmetric learning in pomdpds. In ICML, pages 11013-11023, 2021.
[49] 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.
[50] Daan Wierstra, Alexander Foerster, Jan Peters, and Juergen Schmidhuber. Solving deep memory POMDPs with recurrent policy gradients. In ICANN, pages 697-706, 2007.
[51] 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.
[52] Wenshuai Zhao, Jorge Pena 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.