SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yi Ma
College of Intelligence and Computing
mayi@tju.edu.cn
Hongyao Tang*
tang.hongyao@mila.quebec
Dong Li
Noah's Ark Lab, Huawei Technology dongleecsu@gmail.com
Zhaopeng Meng
mengzp@tju.edu.cn
Offline Reinforcement Learning (RL) aims to address the challenge of distribution shift between the dataset and the learned policy, where the value of out-of-distribution (OOD) data may be erroneously estimated due to overgeneralization. It has been observed that a considerable portion of the benefits derived from the conservative terms designed by existing offline RL approaches originates from their impact on the learned representation. This observation prompts us to scrutinize the learning dynamics of offline RL, formalize the process of generalization, and delve into the prevalent overgeneralization issue in offline RL. We then investigate the potential to rein the generalization from the representation perspective to enhance offline RL. Finally, we present Representation Distinction (RD), an innovative plug-in method for improving offline RL algorithm performance by explicitly differentiating between the representations of in-sample and OOD state-action pairs generated by the learning policy. Considering scenarios in which the learning policy mirrors the behavioral policy and similar samples may be erroneously distinguished, we suggest a dynamic adjustment mechanism for RD based on an OOD data generator to prevent data representation collapse and further enhance policy performance. We demonstrate the efficacy of our approach by applying RD to designed backbone algorithms and widely-used offline RL algorithms. The proposed RD method significantly improves their performance across various continuous control tasks on D4RL datasets, surpassing several state-of-the-art offline RL algorithms.
Reinforcement learning (RL) is a machine learning paradigm centered on training intelligent agents to make decisions through interactions with an environment. In this learning process, an agent develops a policy—a mapping from states to actions—via trial and error, with the goal of maximizing cumulative rewards over time. However, acquiring interaction data in many real-world applications can be expensive, time-consuming, or even hazardous. This spurred the advancement of offline Reinforcement Learning, which has garnered considerable interest in the research community due to its ability to learn effective policies from pre-existing data collected by unknown behavioral policy
without necessitating online interactions with the environment. Offline RL can result in substantial savings in resources, time, and the risk associated with online exploration.
Despite its appealing potential, offline RL encounters challenges in learning optimal policies from limited and suboptimal datasets. The primary challenge in offline RL is the distribution shift between the dataset and the learned policy. As the agent updates its policy, it may come across out-of-distribution (OOD) state-action pairs absent from the limited support of the fixed offline dataset [1]. Conventional RL algorithms, such as SAC [2] or TD3 [3], could yield overly optimistic Q-value estimations for these unseen state-action pairs due to overgeneralization, leading to catastrophic performance. Existing methods in offline RL seek to address this challenge by incorporating techniques like policy constraint [4, 5, 6, 7], conservative value estimates, [8, 9], and uncertainty estimation [10, 11, 12].
Recently, [13] elucidated that the ability of existing offline reinforcement learning (RL) methods to handle OOD actions is significantly attributed to their impact on the learned representations of state-action pairs. The learned representation can capture the underlying structure and essential aspects of the state-action space, enabling the agent to identify similarities and patterns across various state-action pairs, thus improving agent's learning efficacy and generalization capacity. Although there are a few works [14] also study representation learning in offline RL, they focus on how different kinds of pretrained representation matter in downstream policy learning, neglecting the understanding of the essential interplay between representation and policy during the offline co-learning process.
In this paper, we first present a view dubbed Backup-Generalization Cycle to gain a tangible comprehension of how a typical offline RL algorithm learns. We subsequently formalize how generalization happens, and discuss the overgeneralization issue prevalent in offline RL. Finally, we introduce a novel method called Representation Distinction (RD) to mitigate overgeneralization from in-sample state-action pairs to OOD ones. Specifically, we distinctly differentiate the data sourced from the dataset and the data generated by the learning policy from the representation perspective. In cases where the learning policy's performance aligns with the behavioral policy and the generated data largely overlaps with the dataset, we devise an OOD data generator to produce data with lower Q-values than the current policy. We then gradually shift our focus from differentiating in-sample data and data generated by the learning policy to differentiating data generated by the learning policy and data generated by the OOD actor in the representation space, preventing similar samples from being erroneously distinguished.
To demonstrate the efficacy of our proposed RD method in enhancing offline RL policy performance by reining the generalization, we apply it to two backbone algorithms, which are modified versions of the original SAC and TD3 algorithms by incorporating an ensemble of critics and an uncertainty regularization term in the policy update process. By integrating RD into these simple algorithms, our proposed method RD can improve the performance of the backbone agents on D4RL datasets across various continuous control tasks and outperforms several state-of-the-art offline RL algorithms. We also apply RD on two widely-used baselines TD3BC and CQL and improve their performance significantly. In-depth analysis based on visualization and statistical results that demonstrates the efficacy of RD in obtaining more differentiable representation compared to its variants is also provided.
The main contributions of this work can be summarized as follows:
- We introduce a view called Backup-Generalization Cycle to foster an understanding of typical offline value function learning, and highlights the necessity of reining generalization to enhance offline RL.
- We proposed a practical plug-in method Representation Distinction (RD) from the perspective of representation to enhance the performance of offline RL methods by inhibiting overgeneralization among state-action pairs sourced from different distributions.
- We evaluate the effectiveness of RD by applying it to designed backbone algorithms and existing widely-used algorithms and enhance their performance significantly.
Offline RL Offline RL algorithms strive to train RL agents using pre-collected datasets. Nevertheless, the distribution shift between the behavior policy and the policy being learned may lead to issues, as OOD actions are sampled from the learned policy and incorporated into the learned
critic. To address this challenge, various approaches have been proposed. Some earlier methods aim to constrain the learned policy to remain close to the behavior policy, which can be accomplished through explicit policy regularization [4, 5, 6], implicit policy constraints [15, 16, 17, 18, 19], or by employing auxiliary behavioral cloning losses [7]. Alternative approaches penalize the Q-value of OOD actions to discourage their selection [8, 9, 10, 11, 12]. Moreover, model-based methods that train with conservative penalties have been suggested [20, 21, 22, 23, 24, 25].
Representation in Offline RL. Prior research has aimed to analyze various aspects of the representations induced by TD-based methods using function approximation, predominantly in the standard online RL setting [26, 27, 28, 29, 30]. More recently, this line of inquiry has emerged in the offline RL setting [31, 32, 33, 34, 13]. [31] investigates which representations can result in stable convergence of TD in a linear setting. [32, 33] explore the learning dynamics of Q-learning in an overparameterized setting, observing excessively low-rank and aliased feature representations at the fixed points identified by TD-learning. [34] studies the extent of impact of different interventions on the causal link between the effective rank and offline RL agent performance and points out that there is no strong relationship between them. [13] highlights that a substantial portion of the benefits of existing offline RL approaches, which aim to avoid OOD actions, actually originates from their effects on the learned representations. The authors also identify specific metrics that facilitate effective evaluation of the quality of the learned representation. [14] pretrain the representation using different auxiliary loss and then fix the representation to apply it to downstream policy learning. Nevertheless, a thorough examination of the explicit representation distinction between in-sample and OOD state-action pairs in the context of offline RL is lacking. Our approach promotes orthogonality between the representation vectors of in-sample and OOD data, offering a more practical solution for inhibiting overgeneralization, thus enhancing offline RL.
Markov Decision Process (MDP) is a mathematical framework that models decision-making processes in stochastic environments. It is defined by a tuple  $(S, A, P, R, \gamma)$ , where  $S$  is the set of states,  $A$  is the set of actions,  $P$  is the state transition probability function,  $R$  is the reward function, and  $\gamma$  is the discount factor. Q-learning is the major approach for obtaining the optimal  $\pi(s)$ , which learns a Q-value function  $Q(s, a)$  that represents the expected cumulative discounted rewards when starting from the state  $s$  taking the action  $a$  and executing the policy  $\pi$  thereafter. The Q-value function is evaluated by iterating the Bellman operator as  $\mathcal{T}Q(s, a) = \mathbb{E}_{s' \sim P(\cdot | s, a)}\left[r(s, a) + \gamma \mathbb{E}_{a' \sim \pi(s')} Q(s', a')\right]$ . The goal of RL is to find an optimal policy  $\pi(s)$  that maximizes the cumulative discounted rewards.
Offline Reinforcement Learning [1], aims to learn a policy from a fixed dataset of interaction samples, without further interaction with the environment. The dataset  $\mathcal{D}$  consists of transition tuples  $(s,a,r,s^{\prime})$  collected from interactions between one or multiple behavior policies and the environment. The goal is to learn an optimal policy using only this dataset.
To deal with large state-action space, deep RL resorts to deep neural networks for function approximation. Typically, a Q-network parameterized by  $\phi$  can be viewed as  $Q_{\phi}(s,a) = \mathbf{w}^{\top}\Phi (s,a)$ , where  $\Phi (s,a)\in \mathbb{R}^{d}$  is regarded as the representation obtained by the penultimate layer of Q-network while  $\mathbf{w}\in \mathbb{R}^d$  is the linear weight [35, 36, 33, 13]. In conventional online and offline RL learning process, the representation and the linear weight learn together. A few recent works[37, 38, 39, 40] also introduce auxiliary objectives to strengthen or regularize the representation with different purposes. We use this view for the presentation of our analysis and proposed method in the following sections.
4 Reining Generalization in Offline RL
In this section, we focus on the learning dynamics of offline RL. First, we present a view called Backup-Generalization Cycle (Sec. 4.1) to gain a tangible comprehension of how a typical offline RL algorithm learns. We then take a further step to formalize how generalization happens, and discuss the overgeneralization issue in offline RL (Sec. 4.2). Finally, we propose the idea of kernel control to address the overgeneralization issue (Sec. 4.3), with the practical implementation deferred to Sec. 5.
Figure 1: An illustration of Backup-Generalization Cycle.  $s_{\square}, s_{\circ}, s_{\triangle}, s_{\star}$  denotes four states in the offline dataset  $\mathcal{D}$ . Backups and generalizations are denoted by blue and red arrows respectively.
4.1 Backup-Generalization Cycle in Offline RL
Offline RL learns from a static offline dataset with no access to online interaction over the course of policy learning. To gain a better understanding of the intricacies of offline deep RL, for the first, we introduce a view called Backup-Generalization Cycle. This view, as depicted in Fig. 1, fosters an understanding of typical offline value function learning via two key components: Backup and Generalization:
- Typical offline deep RL algorithms employ various forms of backups to learn the  $Q$ -function network. Here we use  $C(\cdot)$  to denote a generic form of the target value such as Clipped Double  $Q$ -Learning [3]. It is notable that only the values of state-action pairs in the offline dataset are directly updated.
- Since the values of all state-action pairs that are not in the dataset will never be directly updated, changes to the values of such  $(s,a) \notin \mathcal{D}$  are solely indirectly instigated by the backups on  $(s,a) \in \mathcal{D}$ , i.e., through generalization. Note that the state-action pairs for the target values of backup are highly likely to be absent from the dataset, i.e.,  $(s',\pi (s')) \notin \mathcal{D}$ , as the learning of current policy  $\pi$ .
Hence, the typical offline learning process of value function can be regarded as a consequence of the complex interplay between backup and generalization. This dynamic interplay forms a cycle: (1) the backups on  $(s,a)\in \mathcal{D}$  consistently influence the values of  $(s,a)\notin \mathcal{D}$  (highly likely including  $(s^{\prime},\pi (s^{\prime})))$ ; (2) the consistently changing  $Q(s^{\prime},\pi (s^{\prime}))$  participates in the backups on  $(s,a)\notin \mathcal{D}$ ; the two kinds of dynamics iterate and twine during the learning process.
There are also some similar opinions mentioned in a few recent works on offline RL [41, 29]. We utilize the Backup-Generalization Cycle to distinctly separate and better analyze these two contributing factors. In the next subsection, we concentrate on the overgeneralization issue with a formal characterization of value generalization among state-action pairs.
4.2 Overgeneralization in Offline RL
According to the Backup-Generalization Cycle presented above, generalization plays a significant role for  $(s,a)\notin \mathcal{D}$ . In contrast to backup, which is relatively explicit and controllable, generalization is implicit and intricate. In the following of this subsection, we analyze value generalization among state-action pairs during the offline RL learning process.
For a starting case, we consider how the  $Q$  function update caused by typical Temporal-Difference (TD) learning on a single state-action pair  $(s,a)\in \mathcal{D}$  (denoted as  $\phi \rightarrow \phi^{\prime}$ ), affects the  $Q$ -value of an arbitrary state-action pair  $(\bar{s},\bar{a})$ . The post-update parameter  $\phi^\prime$  can be formalized as follows:
$$
\phi^ {\prime} = \phi + \left(\mathcal {T} Q _ {\phi} (s, a) - Q _ {\phi} (s, a)\right) \nabla_ {\phi} Q _ {\phi} (s, a), \tag {1}
$$
where learning rate is omitted for convenience. We further formalize the post-update  $Q$ -value of  $(\bar{s},\bar{a})$  by Taylor expansion at the pre-update parameter  $\phi$ :
$$
Q _ {\phi^ {\prime}} (\bar {s}, \bar {a}) = Q _ {\phi} (\bar {s}, \bar {a}) + \nabla_ {\phi} Q _ {\phi} (\bar {s}, \bar {a}) ^ {\top} \left(\phi^ {\prime} - \phi\right) + \mathcal {O} \left(\left\| \phi^ {\prime} - \phi \right\| ^ {2}\right), \tag {2}
$$
Now, we can characterize the generalization from  $(s,a)$  to  $(\bar{s},\bar{a})$  by plugging Eq.1 into Eq.2:
$$
Q _ {\phi^ {\prime}} (\bar {s}, \bar {a}) = Q _ {\phi} (\bar {s}, \bar {a}) + k _ {\phi} (\bar {s}, \bar {a}, s, a) \left(\mathcal {T} Q _ {\phi} (s, a) - Q _ {\phi} (s, a)\right) + \mathcal {O} \left(\| \phi^ {\prime} - \phi \| ^ {2}\right) \tag {3}
$$
Figure 2: A conceptual illustration of Reining Generalization via Kernel Control.
where  $k_{\phi}(\bar{s}, \bar{a}, s, a) \doteq \nabla_{\phi}Q_{\phi}(\bar{s}, \bar{a})^{\top}\nabla_{\phi}Q_{\phi}(s, a)$  and it is called Neural Tangent Kernel in [27]. Eq.3 elucidates the change of the  $Q$ -value of any state-action pair  $(\bar{s}, \bar{a})$  can be mainly characterized by the kernel  $k_{\phi}(\bar{s}, \bar{a}, s, a)$  and the TD-error  $\mathcal{T}Q_{\phi}(s, a) - Q_{\phi}(s, a)$ . Apparently, the kernel controls the extent of generalization: for small values of  $k_{\phi}(\bar{s}, \bar{a}, s, a)$ , the generalization effect is minor and the classic tabular learning (i.e., no generalization) can be viewed as a special case when  $k_{\phi}$  and the higher-order term are 0; for large values of  $k_{\phi}(\bar{s}, \bar{a}, s, a)$ , the generalization effect is prominent and  $Q_{\phi'}(\bar{s}, \bar{a})$  can be largely changed in either the same or the opposite direction of TD-error. Intriguingly, Eq. 3 indicates that we can control the generalization by mainly adjusting the kernel  $k_{\phi}(\bar{s}, \bar{a}, s, a)$ .
It is widely deemed that the crux of offline RL is the overestimation of OOD data, i.e.,  $(s,a)\notin \mathcal{D}$ . Recall the Backup-Generalization Cycle in Fig. 1. The overestimation propagates directly via backups and further spreads via generalization. Most existing works on offline RL focus on eliminating or suppressing overestimation by conservative value estimation [8, 10] or distribution constraint [4, 5], leading to various forms of  $C(\cdot)$ . These works can be categorized as overestimation-control backups. However, since the source of overestimation is improper generalization, i.e., overgeneralization issue, we argue that it is inadequate to only cut the downstream while leave the upstream uncontrolled. Naturally, the generalization should also be carefully controlled as a preposed complement to the overestimation-control backups. Therefore, we take a further step and propose a novel method to control the generalization in the next subsection.
4.3 Reining Generalization via Kernel Control
Since the control of generalization (i.e., the upstream) is neglected in most previous offline RL works, we propose a novel method to address the overgeneralization issue by controlling the kernel in Eq. 3. The key idea is a two-stage generalization control: Policy-Dataset Generalization Inhibition and Policy-OOD Generalization Inhibition. Fig. 2 illustrates the concept of the two stages and an intermediate scene between them, which are detailed below.
As illustrated in the left of Fig. 1, the first stage of generalization control is to inhibit the generalization from behavior policy distribution (or dataset) to learning policy distribution, i.e., suppress  $Q(s,a)\xrightarrow{\text{Generalization}}Q(s,\pi(s))$  for  $(s,a)\in \mathcal{D}$ . Such kind of generalization is called extrapolation in the offline RL literature [1]. Extrapolation is a main component of generalization as its consequence  $Q(s,\pi(s))$  consistently participates in the backups. Thus, overestimation caused by improper extrapolation severely handicaps the learning. For a stable and effective learning process, we argue that extrapolation should be inhibited, especially at the early learning stage, where the learning policy is usually largely different with the behavior policy. This can be achieved by encouraging the kernel to 0, i.e.,  $\min_{\phi}|\nabla_{\phi}Q_{\phi}(s,a)^{\top}\nabla_{\phi}Q_{\phi}(s,\pi(s))|$ . Consequently, training on  $(s,a)\in \mathcal{D}$  induces only a minor change to  $Q_{\phi}(s,\pi(s))$ .
With effective suppression of overestimation achieved by the first stage of generalization inhibition, we further consider an intermediate scene where the learning policy evolves and resembles the behavioral policy as the training proceeds, leading to a distribution overlap. In this scene, the effect of the first stage could become over-inhibition. In another word,  $\min_{\phi}|\nabla_{\phi}Q_{\phi}(s,a)^{\top}\nabla_{\phi}Q_{\phi}(s,\pi (s))|$  for  $a\approx \pi (s)$  could lead to a sharp and bumpy  $Q$  -value landscape as shown in the middle of Fig. 2. It adversely impacts the robustness of  $Q$  function, further resulting in a brittle policy. Hence, it's unwise
to impose the Policy-Dataset Generalization Inhibition all along the learning process. For some empirical evidence of this, we observe performance decay in some of our experiments in Table 3 and Fig.4. A natural remedy is to gradually remove the inhibition. However, it leaves value generalization uncontrolled once again, exposing the policy learning under the risk of collapse due to overestimation.
To address the dilemma above, we introduce the second stage, Policy-OOD Generalization Inhibition, as illustrated in the right of Fig. 2. In this stage, we make use of an additional policy  $\pi_{\mathrm{odd}}$  that outputs OOD actions. The OOD action here is defined as an action with a lower  $Q$ -value than the action selected by the learning policy for the same state. The actual effect of the second stage is suppression of the generalization between the learning policy distribution and OOD policy distribution. Such a means of kernel control can be viewed as a special way of increasing the action gap [42]. It improves the robustness of value function learning by discouraging erroneous generalization on suboptimal actions. Formally, this is achieved by  $\min_{\phi}|\nabla_{\phi}Q_{\phi}(s,\pi (s))^{\top}\nabla_{\phi}Q_{\phi}(s,\pi_{\mathrm{odd}}(s))|$ .
Overall, our kernel control method consists of the two stages of generalization inhibition, which effectively controls the generalization over the entire course of learning, providing a favorable upstream for better offline RL.
5 Representation Distinction in Offline RL
In this section, we devise a practical algorithm to rein the generalization in offline RL following the design of two stages of generalization inhibition. A heuristic method to achieve smooth transition among them is also proposed. We term our algorithm of reining generalization as Representation Distinction (RD) and apply them to both existing and designed backbone algorithms.
5.1 Practical Algorithm Design
We first focus on the phase depicted in the left of Fig.2. A natural approach is to directly minimize  $\mathbb{E}_{s,a\sim \mathcal{D}}|k_{\phi}(s,\pi (s),s,a)|$  across the dataset. However, due to the high dimensionality of  $\nabla_{\phi}Q_{\phi}(s,a)$  the direct computation of  $\mathbb{E}_{s,a\sim \mathcal{D}}|k_{\phi}(s,\pi (s),s,a)|$  involves the computation and backpropagation through per-example gradient dot products, which is computationally prohibitive. To mitigate this, we adopt the method proposed by [33] and approximate  $\Delta (\phi)$  with the contribution solely from the last layer parameters, i.e.,  $\mathbb{E}_{s,a\sim \mathcal{D}}|\nabla_{\mathbf{w}}Q_{\phi}(s,\pi (s))^{\top}\nabla_{\mathbf{w}}Q_{\phi}(s,a)|$ . Consequently, we derive the following loss function:
$$
\mathcal {L} _ {1} = \underset {s, a \sim \mathcal {D}} {\mathbb {E}} | \Phi (s, \pi (s)) ^ {\top} \Phi (s, a) | \tag {4}
$$
where  $\Phi(s, a)$  signifies the representation of state-action pairs as mentioned in Section 3. By minimizing  $\mathcal{L}_1$ , we encourage the learned Q-function to yield representations that are as orthogonal as possible between data from dataset and  $\pi$ . In this manner, the representations of  $(s, a)$  and  $(s, \pi(s))$  can be distinctly discerned.
To substantiate the efficacy of inhibiting generalization between in-sample data and data derived from the learning policy  $\pi$ , we conduct rudimentary experiments by directly minimizing  $\mathcal{L}_1$  during the training of the standard off-policy SAC algorithm [2] on both halfcheetah-medium and halfcheetah-medium-expert. The evaluation results depicted in Fig.3 demonstrate that enhancing the representation alone, devoid
Figure 3: Average results of SAC over 5 seeds with and w/o RD.