SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ke Jiang $^{1,2}$ , Jia-yu Yao $^{3}$ , Xiaoyang Tan $^{1,2*}$
<sup>1</sup> College of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics
$^{2}$  MIIT Key Laboratory of Pattern Analysis and Machine Intelligence
<sup>3</sup> School of Electronic and Computer Engineering, Peking University
ke_jiang@nuaa.edu.cn, jiayu_yao@pku.edu.cn, x.tan@nuaa.edu.cn
We deal with the state distributional shift problem commonly encountered in offline reinforcement learning during test, where the agent tends to take unreliable actions at out-of-sample (unseen) states. Our idea is to encourage the agent to follow the so called state recovery principle when taking actions, i.e., besides long-term return, the immediate consequences of the current action should also be taken into account and those capable of recovering the state distribution of the behavior policy are preferred. For this purpose, an inverse dynamics model is learned and employed to guide the state recovery behavior of the new policy. Theoretically, we show that the proposed method helps aligning the transited state distribution of the new policy with the offline dataset at out-of-sample states, without the need of explicitly predicting the transited state distribution, which is usually difficult in high-dimensional and complicated environments. The effectiveness and feasibility of the proposed method is demonstrated with the state-of-the-art performance on the general offline RL benchmarks.
Reinforcement learning has made significant advances in recent years, but it has to collect experience actively to gain understanding of the underlying environments. However, such online interaction is not always practical, due to either the potential high cost of data collection procedure or its possible dangerous consequences in applications as autonomous driving or healthcare. To address these issues, offline reinforcement learning aims to learn a policy from offline datasets without doing any actual interaction with the environments [19, 13, 1].
However, directly deploying online RL algorithms, such as Deep Deterministic Policy Gradient [18], to learn the new policy from the offline dataset without proper constraints would highly likely suffer from action distributional shift due to the change in the actions generated by the new policy. This would result in the so called extrapolation error [8], i.e., the TD target could be wrongly estimated when querying those out-of-distribution (OOD) actions generated by the new policy. To address these issues, methods like Conservative Q-Learning (CQL) [16], TD3+BC [7] and Implicit Q-Learning (IQL) [14], treat the OOD actions as counterfactual queries and try to avoid performing such queries completely during learning, hence suppressing the Q-function extrapolation error. However, such pessimism for out-of-sample data could be too restricted and sample inefficient, as not all out-of-sample(unseen) states are not generalizable [20].
To effectively generalize to out-of-sample or even OOD states, it is necessary to constrain the behavior of agent at those states, otherwise the policy extrapolation error that occurs at test stage may drive the
Figure 1: An example of CQL agent's failure for the accumulative error due to state distributional shift on a Walker2d robotic agent. The interval between every two images is ten steps.
agent's transited state distribution away from the offline dataset, referred as the state distributional shift problem - a problem has not been paid enough attention as most previous works focus on addressing the issues of OOD actions during training [16, 7]. The problem of state distributional shift is likely to arise in those extensive and non-stationary environments, where it is impossible for the agent to cover the entire state space during training. A brief example is illustrated in Figure 1, where a robotic agent, trained using offline RL frameworks like CQL, is knocked by some external force unintentionally and falls into an out-of-sample state during test, which leads to failure finally due to the unreliable actions taken at this out-of-sample(unseen) state [3] and the error confounding effects.
To deal with the above problem, a natural idea is to teach the agent to recover from these out-of-sample states to their familiar regions, i.e., the demonstration of offline datasets. For this purpose, we propose the Out-of-sample State Recovery (OSR) method to implement this state recovery principle. Our idea is to estimate the actions whose consequence would be in the support region of the offline dataset using inverse dynamics model (IDM) [2]. The estimated IDM can be interpreted as an extended policy with pre-defined consequences, hence being suitable for guiding the learning of the new policy. Theoretically, we show that the proposed OSR makes the transited state distribution of the learnt policy at out-of-sample states align well with the offline dataset, without the need of building any forward dynamics models. We also propose a modified OSR-v that suppresses the probability of selecting actions according to the risk of OOD consequences through imposing extra value-constraints when decision making. We experimentally<sup>2</sup> demonstrated that the proposed OSR outperforms several closely related state-of-the-art methods in offline AntMaze/MuJoCo control suites with various settings.
In what follows, after an introduction and a review of related works, Section 3 provides a concise introduction to the background of offline RL and our work. The OSR method is presented in detail with theoretical analysis of its effectiveness in Section 4.2. The variant of OSR, called OSR-v, is then introduced with theoretical analysis of its advantages in Section 4.3. In Section 4.4, the practical implementation of both methods is described, including the loss function used. Experimental results are presented in Section 5 to evaluate the effectiveness of both methods under various settings. Finally, the paper concludes with a summary of the findings and contributions.
Robust offline reinforcement learning. Robustness is critical for offline RL methods to implement in real-world applications. Recent work RORL [25] adds smoothing term to relax the conservatism of algorithms like Conservative Q-Learning, making the agent generalize to out-of-sample states. ROMI [23] introduce reverse mechanism to build links between isolated data clusters based on out-of-sample data. However, it is still possible that the agent would drift away from the dataset as the results of the well-known error compounding effects due to the myopic of the new policy, referred as state distributional shift problem, which is more important in practical applications. For example, in healthcare [12], it is critical for the agent to make decisions with safe consequences, i.e., generating reliable trajectories. Theoretically, it is shown that control the divergence of the learnt policy from the behavior policy helps to bound the state distributional shift [22], but the bound is actually very loose [15]. Recently model-based State Deviation Correction (SDC) [27] builds a dynamics model and a transition model to guide the agent to generate in-distribution trajectories, through which to enhance the agent's safety and controllability. However, building a high-capacity dynamics model is not always practical in complicated applications, highlighting the necessity to constraint the agent's behavior without constructing any forward model to predict the high-dimensional observations.
2Our code is available at https://github.com/Jack10843/OSR
Inverse dynamics model. An inverse dynamics model  $I(a|s', s)$  predicts the action distribution that explains the transition between a given pair of states. Inverse dynamics models have been applied to improving generalization to real-world problems [6], Markov representation learning [2], defining intrinsic rewards for exploration [5]. In our work, the inverse dynamics model is interpreted as a consequence-constraint policy that generates action distributions with predictable consequences, which in turn plays the role of supervision to the new policy.
A reinforcement learning problem is usually modeled as a Markov Decision Process (MDP), which can be represented by a tuple of the form  $(S, A, P, R, \gamma)$ , where  $S$  is the state space,  $A$  is the action space,  $P$  is the transition probability matrix,  $R$  and  $\gamma$  are the reward function and the discount factor. A policy is defined as  $\pi: S \to A$  and trained to maximize the expected cumulative discounted reward in the MDP:
$$
\max  _ {\pi} \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma R \left(s _ {t}, \pi \left(a _ {t} \mid s _ {t}\right)\right) \right] \tag {1}
$$
In general, we define a Q-value function  $Q^{\pi}(s,a) = \mathbb{E}[\sum_{t=0}^{\infty}\gamma R(s_t,\pi(a_t|s_t))|s,a]$  to represent the expected cumulative rewards. Q-learning is a classic method that trains the Q-value function by minimizing the Bellman error over  $Q$  [24]. In the setting of continuous action space, Q-learning methods use exact or an approximate maximization scheme, such as CEM [11] to recover the greedy policy, as follows,
$$
\begin{array}{l} Q \leftarrow \arg \min  _ {Q} \mathbb {E} \left[ R (s, a) + \gamma \mathbb {E} _ {a ^ {\prime} \sim \pi (\cdot | s ^ {\prime})} Q \left(s ^ {\prime}, a ^ {\prime}\right) - Q (s, a) \right] ^ {2} \\ \pi \leftarrow \arg \max  _ {\pi} \mathbb {E} _ {s} \mathbb {E} _ {a \sim \pi (\cdot | s)} Q (s, a) \tag {2} \\ \end{array}
$$
In offline setting, Q-Learning algorithms learn a Q-value function  $Q^{\pi}(s,a)$  and a policy  $\pi$  from a dataset  $\mathcal{D}$ , which is collected by a behavior policy  $\pi_{\beta}$ . Since there is always an action distributional shift of the new policy  $\pi$  and the behavior policy  $\pi_{\beta}$ , this basic recipe fails to estimate the Q-values for OOD state-action pairs. Conservative Q-Learning (CQL), as a representative OOD-constraint offline RL algorithm, tries to underestimate the Q-values for OOD state-action pairs to prevent the agent from extrapolation error [16]. The CQL term is as follows,
$$
\min  _ {Q} \left[ \mathbb {E} _ {s \sim \mathcal {D}, a \sim \pi (\cdot | s)} Q (s, a) - \mathbb {E} _ {s, a \sim \mathcal {D}} Q (s, a) \right] \tag {3}
$$
However, at test stage, there still exists a static state distributional shift of the new policy  $\pi$  and the behavior policy  $\pi_{\beta}$ , which may accumulate the discrepancy between the agent and demonstration of the dataset. To migrate this, State Deviation Correction (SDC) [27] aims to train a policy choosing actions whose visited states are as closer to the dataset as possible. Specifically, the SDC trains a dynamics model  $M$  and a transition model  $U$ , to optimize the new policy, as follows,
$$
\min  _ {\pi} \lambda \cdot \mathbb {E} _ {s \sim \mathcal {D}} D \left(M (\cdot | \hat {s}, \pi (\cdot | \hat {s})), U (\cdot | s)\right) \tag {4}
$$
where  $\hat{s}$  is a perturbed version of the original state  $s$ ,  $\lambda > 0$  is the weight for the SDC regularization, and  $D$  is a distance measure, which is maximum mean discrepancy (MMD) in [27].
4 Out-of-sample state recovery using inverse dynamics
In this section, we first describe our Out-of-sample State Recovery (OSR) method in detail, as a policy regularization onto the actor loss, which is introduced in detail in section 4.2. Then we also propose a variant of OSR, OSR via value-constraints (OSR-v), for the purposes to reduce the hyperparameters and test the feasibility of practising state recovery principle in different modes, in section 4.3. Finally, we specifically describe how to implement OSR and OSR-v in practice in section 4.4.
4.1 Noise injection onto dataset
It is well-known that noise injection is helpful to address the covariate shift problem in deep learning [17]. While this technique is adopted in [27] to simulate the OOD states, in this work we interpret
the states perturbed with linear Gaussian noise as counter-examples on how to recover from those states into in-sample states. Given an offline dataset  $\mathcal{D}$ , which consists of quadruples  $(s, a, r, s')$ , we first perform data augmentation onto it to form a mixed dataset.
Specifically, given a quadruple  $(s,a,r,s')$ , we perturb the state  $s$  to obtain a noisy state,
$$
\hat {s} = s + \beta \cdot \epsilon \tag {5}
$$
where  $\epsilon$  is sampled from the standard Gaussian distribution  $\mathcal{N}(0,1)$  and  $\beta$  is usually set to be a small constant. Besides, we could also utilize more complex methods to perturb the states, such as the adversarial attacks in [25]. In this way, we obtain a set of perturbed quadruples  $(\hat{s},a,r,s^{\prime})$ , and group them into a new perturbed dataset  $\tilde{D}$ . Note that we do not perturb the next state  $s^\prime$ , to preserve the destination we wish the agent to jump from  $\hat{s}$ .
Finally we combine the two datasets  $\mathcal{D}$  and  $\tilde{\mathcal{D}}$  to form a new dataset  $\mathcal{D}_{tot}$ , i.e.,  $\mathcal{D}_{tot} = \mathcal{D} + \tilde{\mathcal{D}}$ , where each element in  $D_{tot}$  is denoted as  $(\tilde{s},a,r,s')$ . Next we describe how to train a inverse dynamics model to learn how to safely perform transition from  $\{\tilde{s}\}$  to  $s'$ .
4.2 Learning to recover from out-of-sample states via policy-constraints
A inverse dynamics model (IDM), denoted as  $I^{\pi_{\beta}}(a|s,s')$ , is defined in terms of the behavior policy  $\pi_{\beta}$ , dynamics model  $P(s'|s,a)$  and the transition function  $P(s'|s,\pi_{\beta})$  via Bayes' theorem, as follows,
$$
I ^ {\pi_ {\beta}} (a | s, s ^ {\prime}) = \frac {\pi_ {\beta} (a | s) P \left(s ^ {\prime} \mid s , a\right)}{P \left(s ^ {\prime} \mid s , \pi_ {\beta}\right)} \tag {6}
$$
where  $P(s'|s, \pi_{\beta}) = \sum_{a \in A} P(s'|s, a) \pi_{\beta}(a|s)$ . Note that to learn a IDM model, we don't have to estimate the dynamic model of the environment but only needs to treat this task as a usual function approximation problem using a neural network. In particular, using the quadruples  $(\tilde{s}, a, r, s')$  in  $\mathcal{D}_{tot}$ , the needed IDM  $I^{\pi_{\beta}}(a|s, s')$  can be estimated using a probabilistic regression model with  $\tilde{s}, s'$  as input and action  $a$  as output.
With the estimated IDM model available, we interpreted it as an extended policy with desired immediate consequence known, illustrating how to recover from out-of-samples into in-sample(safe) states, then use it to guide the learning of the new policy. In this work, we use the Kullback-Leibler(KL) divergence to measure the divergence between two distributions. In particular, we minimize the KL divergence between  $I^{\pi_{\beta}}(a|\tilde{s},s')$  and the new policy  $\pi (a|\tilde{s})$ , as follows,
$$
\min  _ {\pi} \mathbb {E} _ {\tilde {s} \sim \mathcal {D} _ {t o t}} K L \left(\mathbb {E} _ {s ^ {\prime} \sim P \left(s ^ {\prime} \mid s, \pi_ {\beta}\right)} I ^ {\pi_ {\beta}} \left(a \mid \tilde {s}, s ^ {\prime}\right) \| \pi (a \mid \tilde {s})\right) \tag {7}
$$
where the  $\tilde{s}$  is sampled from the mixed dataset  $\mathcal{D}_{tot}$  and  $s^{\prime}$  is its in-sample consecutive state. (7) is the OSR term aiming to play the same role of traditional model-based SDC, but without modeling the high-dimensional observations. Obviously, (7) utilizes a forward KL divergence, in which the difference between  $I^{\pi_{\beta}}(a|\tilde{s},s^{\prime})$  and  $\pi (a|\tilde{s})$  is weighted by  $I^{\pi_{\beta}}(a|\tilde{s},s^{\prime})$ . Minimizing this forward KL divergence is equivalent to maximizing the policy likelihood as,  $\max_{\pi}\mathbb{E}_{\tilde{s}\sim \mathcal{D}_{tot,s^{\prime}}\sim P(s^{\prime}|s,\pi_{\beta}),a\sim I^{\pi_{\beta}}(a|\tilde{s},s^{\prime})}[\log \pi (a|\tilde{s})]$ , which is easy to achieve by the definition of KL divergence. In this manner, the learnt policy  $\pi$  is able to recover the average behavior of the IDM  $E_{s^{\prime}\sim P(s^{\prime}|s,\pi_{\beta})}I^{\pi_{\beta}}(a|\tilde{s},s^{\prime})$  with greatest probability.
To further explain the feasibility of OSR term, Theorem 1 shows that the OSR term aligns the transited state distribution of the learnt policy with the offline dataset at out-of-sample states in the setting of bounded action space, which is enough to meet the requirements of most RL environments. First, two assumptions are given to serve Theorem 1 in Appendix A, where the assumptions give the continuity and positivity for the transition function of  $\pi_{\beta}$ , and the positivity for the new policy we train. Then,
Theorem 1. Given two consecutive state  $s$  and  $s'$ , the out-of-sample state  $\hat{s}$  within the  $\epsilon$ -neighbourhood of  $s$  and the behavior policy  $\pi_{\beta}$ . Given an inverse dynamics model  $I^{\pi_{\beta}}(a|s,s')$ . In bounded action space setting, the following two optimization formulas are equivalent,
$$
\min  _ {\pi} K L \left(\mathbb {E} _ {s ^ {\prime} \sim P (s ^ {\prime} | s, \pi_ {\beta})} I ^ {\pi_ {\beta}} (a | \hat {s}, s ^ {\prime}) \| \pi (a | \hat {s})\right) \Longleftrightarrow \min  _ {\pi} K L \left(P (s ^ {\prime} | s, \pi_ {\beta}) \| P (s ^ {\prime} | \hat {s}, \pi)\right)
$$
, that is, the two optimization formulas achieve the same policy.
The proof of Theorem 1 is conducted in Appendix B.1. It theoretically guarantees that OSR have the effect that the learnt policy is regularized to recover the transited state distribution  $P(s'|\hat{s},\pi)$  at out-of-sample state  $\hat{s}$  aligned well with the transited state distribution  $P(s'|s,\pi_{\beta})$  of the behavior policy  $\pi_{\beta}$  at the corresponding in-sample state  $s$ . For the fact that the offline dataset is generated via the behavior policy  $\pi_{\beta}$ , then  $P(s'|s,\pi_{\beta})$  is the in-distributional transited state distribution of the offline dataset. Therefore, we remark that OSR helps aligning the new policy's transited state distribution  $P(s'|\hat{s},\pi)$  at out-of-sample state  $\hat{s}$  with the offline dataset. Interestingly, we find that the term  $\min_{\pi}KL\big(P(s'|s,\pi_{\beta})\big||P(s'|\hat{s},\pi)\big)$  could also be termed as a KL-divergence version of SDC introduced in (4).
4.3 Value-constraint out-of-sample state recovery: a variant
In this section, we propose Out-of-sample State Recovery via value-constraints (OSR-v) as a variant of OSR proposed before. The proposed OSR-v penalizes the Q-values of those actions that prefer transiting to OOD states at out-of-sample states, which drives the agent able to choose the actions that transits to in-sample states, recovering from out-of-sample states. In particular, given two consecutive states  $\tilde{s}$  and  $s^{\prime}$  sampled from the mixed dataset  $\mathcal{D}_{tot}$  constructed before, we overestimate the Q-values of actions that transits to in-sample states, i.e., actions generated by the IDM  $I^{\pi_{\beta}}(\cdot |\tilde{s},s^{\prime})$  while suppressing the Q-values of actions with unknown consequences, i.e., actions generated by the new policy  $\pi$ , as follows,
$$
\min  _ {Q} \mathbb {E} _ {\tilde {s}, s ^ {\prime} \sim \mathcal {D} _ {t o t}} \left(\mathbb {E} _ {\hat {a} \sim \pi (\cdot | \tilde {s})} Q (\tilde {s}, \hat {a}) - \mathbb {E} _ {\hat {a} _ {t a r} \sim I ^ {\pi_ {\beta}} (\cdot | \tilde {s}, s ^ {\prime})} Q (\tilde {s}, \hat {a} _ {t a r})\right) \tag {8}
$$
where  $\hat{a}_{tar}$  refers to the action sampled from IDM, which is referred as target action.
In practice, we observe that OSR-v enable the agent to follow the state recovery principle as well. In order to rigorously explain this phenomenon, we give Theorem 2 under the assumptions mentioned in Appendix A to show that performing OSR-v is equivalent to underestimating the value of OOD states while overestimating the value of in-sample states, which guides the agent prefer to transit to those in-sample(safe) states, avoiding error accumulated via state distributional shift.
Theorem 2. Given two consecutive state  $s$  and  $s'$ , the out-of-sample state  $\hat{s}$  within the  $\epsilon$ -neighbourhood of  $s$  and the behavior policy  $\pi_{\beta}$ .  $\pi$  is the new policy. Define  $Q(s, a)$  as a proximal  $Q$ -value function and  $V(s)$  as a proximal value function. Given an inverse dynamics model  $I^{\pi_{\beta}}(a|s, s')$ . Then the following two optimization formulas are approximately equivalent:
$$
\min  _ {Q} \mathbb {E} _ {\hat {a} \sim \pi (\cdot | \hat {s})} Q (\hat {s}, \hat {a}) - \mathbb {E} _ {\hat {a} _ {t a r} \sim I ^ {\pi_ {\beta}} (\cdot | \hat {s}, s ^ {\prime})} Q (\hat {s}, \hat {a} _ {t a r}) \tag {9}
$$
$$
\stackrel {\Leftrightarrow} {\min } _ {V} \mathbb {E} _ {s ^ {\prime} \sim P \left(s ^ {\prime} \mid \hat {s}, \pi\right)} V \left(s ^ {\prime}\right) - \mathbb {E} _ {s ^ {\prime} \sim P \left(s ^ {\prime} \mid s, \pi_ {\beta}\right)} V \left(s ^ {\prime}\right) \tag {10}
$$
, that is, they achieve the same policy approximately. Then the difference of the two optimization formulas is no larger than  $2\delta \sum_{a}\pi_{\beta}(a|\hat{s})Q(\hat{s},a)$ .
The proof of Theorem 2 is conducted in Appendix B.4. Theorem 2 demonstrates the approximate equivalence of (9) and (10). There is actually a little gap between (9) and (10) which is caused by the existence of the continuous gap  $\delta$ , so when  $\delta$  is large, (9) and (10) are quite different. However, if we assume the infinite norm of reward function  $\|R\|_{\infty} \leq 1$ , then  $2\delta \sum_{a} \pi_{\beta}(a|\hat{s}) Q(\hat{s}, a) \leq \frac{2\delta}{1 - \gamma}$ , where  $\gamma$  is the discount factor, and by the conservative selection of the magnitude of the noise  $\epsilon$ , the  $\delta$  would be a quite tiny value in practice.
The  $s'$  sampled from  $P(s'|s, \pi_{\beta})$ , the state distribution aligned with the offline dataset  $\mathcal{D}$ , is in-sample, then overestimating of such  $s'$  increases the agent's preference to transiting to in-sample states, which is supported by the reasoning presented in Proposition 1, as follows,
Proposition 1. Given an arbitrary state  $s$  and its target state  $s'$ .  $Q$  is an approximal  $Q$ -value function and  $V$  is its value function.  $\pi_Q$  is the learnt policy according to  $Q$ . Let  $\pi_Q(a|s)$  is positive correlated to  $Q(s,a)$ , noted as  $\pi_Q(a|s) \propto Q(s,a)$ , and assume  $\exists a, P(s'|s,a) > 0$ , then  $P(s'|s,\pi_Q) \propto V(s')$ .
The proof and more discussion of Proposition 1 is conducted in Appendix B.2. For the fact that out-of-sample(noisy) state  $\hat{s}$  is close to  $s$ , the in-sample prior of  $s'$ , so we assume  $\exists a, P(s'|\hat{s}, a) > 0$  and via Proposition 1,  $\mathbb{E}_{s' \sim P(s'|s, \pi_{\beta})} P(s'|\hat{s}, \pi) \propto E_{s' \sim P(s'|s, \pi_{\beta})} V(s')$ . Therefore, the result of Theorem 2 and Proposition 1 guarantees that OSR- $\nu$  can align the transited state distribution of the new policy with the offline dataset, avoiding compounding state distributional shift as well.
4.4 Implementation and algorithm summary
We implement our methods, OSR and OSR-v, introduced in previous sections based on an Conservative Q-Learning(CQL) framework to practice the state recovery principle in the setting of countinous action space, without modeling the dynamics of the environment. To be specific, we would implement OSR in policy-constraint mode and OSR-v in value-constraint mode.
Policy-constraint mode (OSR). The policy-constraint based method, OSR proposed in (7), could be implemented as SR loss  $L_{sr}$  in (11), where the detailed derivation is presented in Appendix B.5.
$$
L _ {s r} = E _ {\tilde {s}, s ^ {\prime} \sim \mathcal {D} _ {t o t}} K L \left(I ^ {\pi_ {\beta}} \left(a | \tilde {s}, s ^ {\prime}\right) \| \pi (a | \tilde {s})\right) \tag {11}
$$