SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Peng Cheng\*1,3, Xianyuan Zhan\*†2,4, Zhihao Wu†1,3, Wenjia Zhang2, Shoucheng Song1,3, Han Wang1,3, Youfang Lin1,3, Li Jiang2
<sup>1</sup> Beijing Jiaotong University, Beijing, China
$^{2}$  Tsinghua University, Beijing, China
<sup>3</sup> Beijing Key Laboratory of Traffic Data Analysis and Mining, Beijing, China
$^{4}$  Shanghai Artificial Intelligence Laboratory, Shanghai, China
pcheng6@126.com, zhanxianyuan@air.tsinghua.edu.cn,
zhwu，yflin@bjtu.edu.cn
Offline reinforcement learning (RL) offers an appealing approach to real-world tasks by learning policies from pre-collected datasets without interacting with the environment. However, the performance of existing offline RL algorithms heavily depends on the scale and state-action space coverage of datasets. Real-world data collection is often expensive and uncontrollable, leading to small and narrowly covered datasets and posing significant challenges for practical deployments of offline RL. In this paper, we provide a new insight that leveraging the fundamental symmetry of system dynamics can substantially enhance offline RL performance under small datasets. Specifically, we propose a Time-reversal symmetry (T-symmetry) enforced Dynamics Model (TDM), which establishes consistency between a pair of forward and reverse latent dynamics. TDM provides both well-behaved representations for small datasets and a new reliability measure for OOD samples based on compliance with the T-symmetry. These can be readily used to construct a new offline RL algorithm (TSRL) with less conservative policy constraints and a reliable latent space data augmentation procedure. Based on extensive experiments, we find TSRL achieves great performance on small benchmark datasets with as few as  $1\%$  of the original samples, which significantly outperforms the recent offline RL algorithms in terms of data efficiency and generalizability. Code is available at: https://github.com/pcheng2/TSRL
The recently emerged offline reinforcement learning (RL) provides a new paradigm to learn policies from pre-collected datasets without the need to interact with the environments [1, 2, 3]. This is particularly desirable for solving practical tasks, as interacting with real-world systems can be costly or risky, and high-fidelity simulators are also hard to build [4]. However, existing offline RL methods have high requirements on the size and quality of offline datasets in order to achieve reasonable performance. When such requirements are not met, these algorithms may suffer from severe performance drop, as illustrated in Figure 1. Current offline RL algorithms are trained and validated on benchmark datasets (e.g., D4RL [5]) that contain millions of transitions for simple tasks. Whereas under realistic settings, it is often impractical or costly to collect such a large amount of data, and the real datasets might only narrowly cover the state-action space. Clearly, learning reliable policies from
small datasets with partial coverage has become one of the most pressing challenges for successful real-world deployments of offline RL.
Unfortunately, sample-efficient design considerations have been largely overlooked in the majority of offline RL studies. Pessimism is universally adopted in existing offline RL methods and various forms of data-related regularizations have been applied to combat the distributional shift and exploitation error accumulation issues [3, 6], such as conservatively restricting policy deviation from the behavioral data [3, 6, 7, 8], regularizing value function on out-of-distribution (OOD) samples [9, 10, 11, 12], learning policy on a pessimistic MDP [4, 13, 14],
Figure 1: Performance of recent offline RL methods and our proposed TSRL on the D4RL-MuJoCo Walker2d medium and expert datasets when reducing the number of samples from 1M (full dataset) to  $10\mathrm{k}$ $(1\%)$ .
or adopting strict in-sample learning [15, 16, 17, 18]. In many of these methods, the full coverage assumption plays an important role in their theoretical performance guarantees [3, 19], which assumes the dataset to contain all state-action pairs in the induced distribution of the policy. Obviously, most of the state-action space will become OOD regions under a small dataset. Applying strict data-related regularizations will inevitably cause severe performance degradation and poor generalization [20]. Consequently, it is important to rethink what is essential in policy learning with small datasets. In other words, what is the fundamental or invariant information that can be used to facilitate policy learning, without being conservatively confined by the limited data?
In this paper, we provide a new insight that exploiting the fundamental symmetries in the system dynamics can substantially enhance the performance of offline RL with small datasets. Specifically, we consider the time-reversal symmetry (also called  $T$ -symmetry), which is one of the most fundamental properties discovered in classical and quantum mechanics [21, 22]. It suggests that the underlying laws of physics should not change under the time-reversal transformation:  $t \rightarrow -t$  [22, 23]. Specifically, we are interested in an extended form of T-symmetry for MDP due to its simplicity and universality in physical systems. In the small dataset setting, enforcing T-symmetry in dynamics model learning and offline RL offers three crucial benefits. First, as T-symmetry describes a fundamental property of a system, hence learning it in principle would not require a large or high-coverage dataset. Second, T-symmetry captures what is essential and invariant in the dynamical system, thus offering great generalizability. Lastly, compliance with T-symmetry also provides an important clue to detecting unreliable or non-generalizable state-action samples.
Based on these intuitions, we develop a physics-informed T-symmetry enforced Dynamics Model (TDM) to learn a well-behaved and generalizable dynamics model with small datasets. TDM enforces the extended T-symmetry between a pair of latent space forward and reverse dynamics sub-models, which are modeled as first-order ordinary differential equation (ODE) systems to extract fundamental dynamics patterns in data. TDM provides both well-behaved representations for small datasets and a new reliability measure for OOD samples based on compliance with the T-symmetry. These can be used to construct a highly data-efficient offline RL algorithm, which we call T-Symmetry regularized offline RL (TSRL). Specifically, TSRL uses the T-symmetry regularized representations learned in TDM to facilitate value function learning. Furthermore, the deviation on latent actions and the consistency with T-symmetry specified in TDM provide another perspective to detect unreliable or non-generalizable samples, which can serve as a new set of policy constraints to replace the highly restrictive OOD regularizations in existing offline RL algorithms. Lastly, a reliable latent space data augmentation scheme based on compliance with the T-symmetry is also applied to further remedy the limited size of training data. With these designs, TSRL performs surprisingly well compared with the state-of-the-art offline RL algorithms on reduced-size D4RL benchmark datasets with even as few as  $1\%$  of the original samples. To the best of the authors' knowledge, this is the first offline RL method that demonstrates promising performance on extremely small datasets.
Offline reinforcement learning. We consider the standard Markov decision process (MDP) setting [24], which is represented as a tuple  $\mathcal{M} = \{S, \mathcal{A}, r, \mathcal{P}, \rho, \gamma\}$ , where  $S$  and  $\mathcal{A}$  are the state and
action spaces,  $r(s,a)$  is a scalar reward function,  $\mathcal{P}$  is the transition dynamics,  $\rho$  is the initial state distribution, and  $\gamma \in (0,1)$  is a discount factor. The objective of RL is to learn a policy  $\pi(a|s)$  by maximizing the expected cumulative discounted return  $\mathbb{E}_{\pi}[\sum_{t=0}^{\infty} \gamma^{t} r(s_{t},a_{t})]$ , which is typically approximated by a value function  $Q(s,a)$  using some function approximators, such as deep neural networks. The Q-function is typically learned by minimizing the squared Bellman error:
$$
Q = \underset {Q} {\arg \min } \mathbb {E} \left[ \left(Q (s, a) - \mathcal {B} ^ {\pi} \hat {Q} (s, a)\right) ^ {2} \right] \tag {1}
$$
where  $\hat{Q}$  denotes a target Q-function, which is a delayed copy of the current Q-function;  $\mathcal{B}^{\pi}$  is the Bellman operator, which is often used as the Bellman evaluation operator  $\mathcal{B}^{\pi}\hat{Q}(s,a) = r(s,a) + \gamma \mathbb{E}_{a'}\sim_{\pi}\hat{Q}(s',a')$  in many RL algorithms.
Under the offline RL setting, we are provided with a fixed dataset  $\mathcal{D} = \{(s_0, a_0, r_0, s_1, \dots)^{(i)}\}_{i=1}^N$  without any chance of further environment interactions. Directly applying standard online RL methods in the offline setting suffers from severe value overestimation, due to counterfactual queries on OOD data and the resulting extrapolation errors [1, 3, 6]. To avoid this issue, a widely used offline RL framework adopts the following behavior regularization scheme which regularizes the divergence between the learned policy  $\pi$  and the behavior policy  $\pi_{\beta}$  of the dataset  $\mathcal{D}$ :
$$
\pi = \underset {\pi} {\arg \max } \mathbb {E} _ {s \sim \mathcal {D}, a \sim \pi (\cdot | s)} [ Q (s, a) - \mathrm {D} (\pi (\cdot | s) \| \pi_ {\beta} (\cdot | s)) ] \tag {2}
$$
where  $\mathrm{D}(\cdot \| \cdot)$  is some divergence measures, which can have either an explicit [3, 8] or implicit form [6, 17, 25]. Although straightforward, existing behavior regularization methods have been shown to be over-conservative [9, 20] due to the restrictive regularization with respect to the behavior policy in data, which may suffer from notable performance drop under small datasets.
Time-reversal symmetry in dynamical systems. Most real-world dynamical systems with state measurement  $\mathbf{x} \in \Omega$  on some phase space  $\Omega$  can be modeled or approximated by the system of non-linear first-order ordinary differential equations (ODEs) as  $\frac{d\mathbf{x}}{dt} = F(\mathbf{x})$ , where  $F$  is some general non-linear, at least  $\mathcal{C}^1$ -differentiable vector-valued function. First-order ODE systems are said to be time-reversal symmetric if there is an invertible transformation  $\Gamma: \Omega \mapsto \Omega$ , that reverses the direction of time [22, 26]:  $d\Gamma(\mathbf{x}) / dt = -F(\Gamma(\mathbf{x}))$ . If we define a time evolution operator  $U_{\Delta t}: \Omega \mapsto \Omega$  as  $U_{\Delta t}: \mathbf{x}(t) \mapsto U_{\Delta t}(\mathbf{x}(t)) = \mathbf{x}(t + \Delta t)$ . Then T-symmetry implies that  $\Gamma \circ U_{\tau} = U_{-\tau} \circ \Gamma$ . In other words, the reversing of the forward time evolution of an arbitrary state should be equal to the backward time evolution of the reversed state.
Extending T-symmetry for more generic MDP settings. In our discrete-time MDP setting, we have  $\mathbf{x} = (s, a)$ . We can slightly abuse the notations and denote  $\dot{s} = \frac{ds}{dt}$  as the time-derivative of the current state  $s$ , which can be approximated as the difference between the next and current states, i.e.,  $\dot{s} = s' - s$ . For a dynamical system that satisfies T-symmetry, it suggests that if we learn a forward dynamics  $F(s, a) = \dot{s}$  and a reverse dynamics  $\tilde{G}(s', a') = -\dot{s}$  as a pair of first-order ODEs, we should have  $F(s, a) = -\tilde{G}(s', a')$ .
However, from a decision-making perspective, it is known that T-symmetry can sometimes be broken by irreversible actions or some special dynamic processes (e.g., frictional force against motion). Hence in this paper, we consider a more generic treatment by leveraging an alternative ODE reverse dynamics model  $G(s', a) = -\dot{s}$  to establish the T-symmetry with the forward dynamics, i.e., enforcing  $F(s, a) = -G(s', a)$ . Note that  $G(s', a)$  is now defined on the next state  $s'$  and the current action  $a$ , rather than the next action  $a'$ , thus is not impacted if the next action is irreversible. This extended T-symmetry provides a more fundamental and almost universally held property in discrete-time MDP systems. Its simplicity and fundamentalness make it an ideal property that we can leverage to construct a well-behaved data-driven dynamics model and a robust offline RL algorithm under small datasets.
3 T-Symmetry Enforced Dynamics Model
In this section, we present the detailed design of TDM, which is capable of learning a more fundamental and T-symmetry preserving dynamics from small datasets. The key ingredients of TDM are to embed a pair of latent forward and reverse dynamics as ODE systems, and further enforce their T-symmetry consistency. This design offers several benefits. First, embedding ODE modeling
Figure 2: Overall architecture of the proposed TDM
of the system helps to extract more fundamental descriptions of how the system dynamics evolve over time in the latent space. Similar ideas have also been explored in Koopman theory [27, 28] and discovering governing equations of nonlinear dynamical systems [29, 30]. The requirement on extended T-symmetry for both forward and reverse dynamics introduces additional consistency regularization, which is helpful to improve model learning stability under limited data. With a regularized and well-behaved dynamics characterization, we can learn more effective state-action representations, which allows removing non-essential elements in the raw data to promote generalization performance. Lastly, compliance with T-symmetry can provide extra information regarding the reliability of OOD samples, which can be particularly useful in downstream offline policy learning. As illustrated in Figure 2, the proposed TDM consists the following components:
Encoder and decoders. TDM implements a state-action encoder  $\phi(s, a) = (z_s, z_a)$  and a pair of decoders  $\psi_s(\cdot, \delta_s), \psi_a(z_a) = a$  that embed the state-action pair  $(s, a)$  into latent representations  $(z_s, z_a)$  and then map them back. Specifically, we require the state decoder  $\psi_s(\cdot, \delta_s)$  to be capable of decoding both  $z_s$  and  $\dot{z}_s$ , where  $\delta_s$  is an indicator to help the decoder to decide the target output, with  $\delta_s = 0$  as decoding  $z_s \to s$  and  $\delta_s = 1$  as decoding  $\dot{z}_s \to \dot{s}$ . The encoder  $\phi$  and decoders  $\psi_s$  and  $\psi_a$  induce the following reconstruction loss term for each state-action pair  $(s, a)$ :
$$
\ell_ {r e c} (s, a) = \left\| s - \psi_ {s} \left(z _ {s}, 0\right) \right\| _ {2} ^ {2} + \left\| a - \psi_ {a} \left(z _ {a}\right) \right\| _ {2} ^ {2} \tag {3}
$$
Latent forward dynamics. Inspired by the prior works that incorporate physics-informed information into dynamical systems modeling [27, 29, 30], we embed a discrete-time first-order ODE system to capture the latent forward dynamics  $f(z_{s},z_{a}) = \dot{z}_{s}$ . Similar to  $\dot{s}$ , we write  $\dot{z}_s = z_{s'} - z_s$  to denote the forward difference of the next and current latent state representations. Note that based on the chain-rule, we have  $\dot{z}_s = \frac{dz_s}{dt} = \frac{\partial z_s}{\partial s}\cdot \frac{ds}{dt} = \nabla_s z_s\cdot \dot{s}$ . To enforce the ODE property, we can introduce the following loss term for  $f$ :
$$
\ell_ {f w d} (s, a, s ^ {\prime}) = \| \left(\nabla_ {s} z _ {s}\right) \dot {s} - \dot {z} _ {s} \| _ {2} ^ {2} = \| \frac {\partial \phi (s , a)}{\partial s} \dot {s} - f (\phi (s, a)) \| _ {2} ^ {2} \tag {4}
$$
Minimizing  $\mathcal{L}_{fwd}$  ensures that the latent forward dynamics  $f$  correctly predicts the forward time evolution of latent states in the dynamical system. We also require the decoder  $\psi_s(\cdot, \delta_{\dot{s}})$  to be able to decode  $\dot{s}$  from  $\dot{z}_s$  to ensure it is compatible with the ODE property, which implies the following loss:
$$
\ell_ {d s} (s, a, s ^ {\prime}) = \| \dot {s} - \psi_ {s} (\dot {z} _ {s}, 1) \| _ {2} ^ {2} = \| \dot {s} - \psi_ {s} (f (\phi (s, a)), 1) \| _ {2} ^ {2} \tag {5}
$$
Latent reverse dynamics. We can further introduce a latent reverse dynamics  $g(z_{s'}, z_a) = -\dot{z}_s$  in the model, which captures the reverse time evolution of the system in the latent space. Similar to the forward dynamics loss  $\mathcal{L}_{fwd}$ , we can write the reverse dynamics loss for  $g$  as:
$$
\ell_ {r v s} (s, a, s ^ {\prime}) = \left\| \left(\nabla_ {s ^ {\prime}} z _ {s ^ {\prime}}\right) (- \dot {s}) - (- \dot {z} _ {s}) \right\| _ {2} ^ {2} = \left\| \frac {\partial \phi \left(s ^ {\prime} , a\right)}{\partial s ^ {\prime}} (- \dot {s}) - g \left(\phi \left(s ^ {\prime}, a\right)\right) \right\| _ {2} ^ {2} \tag {6}
$$
T-symmetry regularization. The above latent forward and reverse dynamics  $f$  and  $g$  are learned to be two models, which may not necessarily satisfy the proposed extended T-symmetry. We can enforce the extended T-symmetry by requiring  $f(z_{s},z_{a}) = -g(z_{s^{\prime}},z_{a})$ . To further couple the learning process of  $f$  and  $g$ , note that  $z_{s^{\prime}} = z_{s} + \dot{z}_{s} = z_{s} + f(z_{s},z_{a})$ , which suggests  $g(z_{s^{\prime}},z_{a}) = g(z_{s} + f(z_{s},z_{a}),z_{a}) = -\dot{z}_{s} = -f(z_{s},z_{a})$ . This implies the following T-symmetry consistency loss:
$$
\ell_ {T - s y m} \left(z _ {s}, z _ {a}\right) = \left\| f \left(z _ {s}, z _ {a}\right) + g \left(z _ {s} + f \left(z _ {s}, z _ {a}\right), z _ {a}\right) \right\| _ {2} ^ {2} \tag {7}
$$
Above instance-wise T-symmetry consistency loss also provides an alternative measure for evaluating the reliability of a data sample. A state-action pair  $(s, a)$  with a large  $\ell_{T - sym}(\phi(s, a))$  implies that this sample may not be well-explained by TDM or consistent with the fundamental symmetry of the system. This can be used to detect unreliable OOD samples in offline policy optimization as well as construct a new latent space data augmentation procedure, which will be discussed in later content.
Final learning objective of TDM. Finally, we can formulate the overall loss function of TDM as:
$$
\mathcal {L} _ {T D M} = \sum_ {(s, a, s ^ {\prime}) \in \mathcal {D}} \left[ \ell_ {r e c} + \ell_ {d s} + \ell_ {f w d} + \ell_ {r v s} + \ell_ {T - s y m} \right] (s, a, s ^ {\prime}) + \lambda_ {L 1} \left[ \mathcal {L} _ {1} (f) + \mathcal {L} _ {1} (g) \right] \tag {8}
$$
where  $\mathcal{L}_1(f)$  and  $\mathcal{L}_1(g)$  are L1-norms of the parameters of  $f$  and  $g$ , and  $\lambda_{L1}$  is a scale parameter. L1 regularization is introduced to encourage learning parsimonious latent dynamics for  $f$  and  $g$ , which helps to improve model generalizability [29, 30].
Note that the proposed TDM is very different from the conventional dynamics models used in model-based RL (MBRL) methods [4, 13, 14, 31, 32, 33, 34, 35]. The dynamics models in MBRL focus on constructing a predictive model to represent the forward transition dynamics of the system. Whereas, TDM is formulated as a reconstruction model with T-symmetry preserving embedded ODE latent dynamics, which aims at explaining and extracting the fundamental dynamics of the system. As a result, TDM can be substantially more well-behaved and robust when learning from small datasets.
4 T-Symmetry Regularized Offline RL
In this section, we discuss how to incorporate the properties of TDM to construct a sample-efficient offline RL algorithm, which we call T-Symmetry regularized offline RL (TSRL). Specifically, we leverage three types of information from TDM to jointly improve offline RL performance under small datasets. First, we use the well-behaved state-action representations provided by TDM to facilitate value function learning. Second, the consistency with the latent dynamics in TDM also provides new forms of policy constraints to penalize unreliable and non-generalizable OOD samples. Finally, compliance with T-symmetry enables a new latent space data augmentation procedure to further enhance the algorithm's performance with limited data.
T-symmetry regularized representation. Representation learning has been shown to be an effective approach to enhancing sample efficiency and generalization in many online and offline RL studies [36, 37, 38, 39, 40]. A notable property of TDM is that the learned latent state-action representations from the encoder  $(z_{s}, z_{a}) = \phi(s, a)$  are compatible with both the latent forward and reverse ODE dynamics  $f$  and  $g$ . This leads to well-regularized and T-symmetry preserving representations that can potentially generalize better on OOD areas under small dataset settings. We can simply use the latent state-action representation  $(z_{s}, z_{a})$  extracted by the encoder  $\phi(s, a)$  of TDM in the value function learning, which gives the following policy evaluation objective:
$$
Q = \underset {Q} {\operatorname {a r g m i n}} \mathbb {E} _ {(s, a, s ^ {\prime}) \sim \mathcal {D}} \left[ \left(r (s, a) + \gamma \hat {Q} \left(\phi \left(s ^ {\prime}, \pi (\cdot | s ^ {\prime})\right)\right) - Q \left(\phi (s, a)\right)\right) ^ {2} \right] \tag {9}
$$
T-symmetry regularized policy constraints. Existing offline RL methods primarily penalize the divergence between the learned policy  $\pi$  and the behavioral data in the original action space, which ignores the underlying manifold structure of actions in the latent space [41] and the system dynamics properties. Moreover, restricting policy optimization only within data-covered regions can be overconservative when the offline dataset is small, which can result in degraded policy performance [20]. In TSRL, we instead consider an alternative regularization scheme, which restricts the deviation on latent actions and the T-symmetry consistency of policy-induced samples, corresponding to the following policy optimization objective:
$$
\underset {\pi} {\operatorname {a r g m a x}} \mathbb {E} _ {(s, a) \sim \mathcal {D}} \left[ \alpha Q \left(\phi (s, \pi (\cdot | s)) - \lambda_ {1} \| z _ {a ^ {\pi}} - z _ {a} \| _ {2} ^ {2} - \lambda_ {2} \ell_ {T - s y m} (\phi (s, \pi (\cdot | s))) \right] \right. \tag {10}
$$
where latent actions  $z_{a}$  and  $z_{a^{\pi}}$  are obtained from  $\phi(s, a)$  and  $\phi(s, \pi(\cdot|s))$  respectively. The second term restricts the latent action  $z_{a^{\pi}}$  of policy  $\pi$  from deviating too much from the latent action  $z_{a}$  in data. The third term regularizes the T-symmetry consistency of policy-induced samples  $(s, \pi(\cdot|s))$ , which is evaluated based on Eq. (7) and the learned TDM.  $\lambda_{1}$  and  $\lambda_{2}$  are weight parameters, which only need to be roughly adjusted to ensure both the regularization terms are in a similar scale
as the first term. We also introduce a normalization term  $\alpha$  on the value function for training stability similar to TD3+BC [8], which is computed based on a training batch  $B$  of samples as  $\alpha_0 / [\sum_{(s,a)\in B}Q(\phi (s,\pi (\cdot |s)))]$ . We set  $\alpha_0 = 2.5$  in all of our experiments without tuning.
Instead of strictly regularizing OOD actions as in existing offline RL algorithms, our policy regularization scheme actually allows some reliable OOD action for policy optimization. As TDM is designed to capture the fundamental and invariant system dynamics patterns in the latent space, if an OOD action has a similar latent representation to some latent actions in data and also agrees with the T-symmetry property in TDM, then we can expect some degree of equivalency between these actions. This leads to more relaxed policy constraints by enabling policy learning and generalization on reliable OOD regions, which is critical for the small dataset setting.