SlowGuess/NeurIPS_2023_OCR huggingface dataset:
T-symmetry consistent latent space data augmentation. It has been shown in previous studies [42, 28, 43] that data augmentation can potentially improve the function approximation of the Q-networks by smoothing out the learned state-action space, hence often lead to more robust policy and better data efficiency. However, existing data augmentation methods in offline RL studies either blindly add random perturbations to states [42] or utilize costly non-linear symmetry transformations, such as Koopman theory [28]. With TDM, we can provide a very simple yet principled data augmentation scheme based on the T-symmetry property. Assuming we add a small perturbation  $\epsilon$  to a latent state  $z_{s}$ , i.e.,  $(z_{s}, z_{a}) \mapsto (z_{s} + \epsilon, z_{a})$ , then the corresponding perturbation  $\epsilon'$  on the next latent state  $z_{s'}$  according to the latent forward dynamics  $\dot{z}_{s} = f(z_{s}, z_{a})$  satisfies:  $z_{s'} + \epsilon' = z_{s} + \epsilon + f(z_{s} + \epsilon, z_{a})$ . On the other hand, by the T-symmetry construction in TDM, we can recover back the current perturbed latent state based on the latent reverse dynamics  $-\dot{z}_{s} = g(z_{s'}, z_{a})$  as:  $z_{s} + \epsilon'' = z_{s'} + \epsilon' + g(z_{s'} + \epsilon', z_{a})$ . Clearly, we should have  $\epsilon = \epsilon''$ , which suggests the following condition:
$$
\epsilon^ {\prime \prime} - \epsilon = f \left(z _ {s} + \epsilon , z _ {a}\right) + g \left(z _ {s} + \epsilon + f \left(z _ {s} + \epsilon , z _ {a}\right), z _ {a}\right) = 0 \tag {11}
$$
This is exactly equivalent to requiring the instance-wise T-symmetry consistency loss (Eq. (7))  $\ell_{T - sym}(z_s + \epsilon ,z_a) = 0$ . Hence we can use T-symmetry consistency loss  $\ell_{T - sym}(\cdot)$  as a reliability measure to filter out unreliable augmented samples  $(z_{s} + \epsilon ,z_{a})$  that are inconsistent with the T-symmetry property of the learned latent dynamics in TDM. In our implementation, we only keep augmented samples that satisfy  $\ell_{T - sym}(z_s + \epsilon ,z_a)\leq h$ , where we consider a non-parametric treatment for threshold  $h$ , by setting it as the  $\tau$ -quantile value of all  $\ell_{T - sym}(\phi (s,a))$  values of  $(s,a)$  in  $\mathcal{D}$  (we choose  $\tau = 50\%$  or  $70\%$  in our experiments). This ensures that the augmented samples at least maintain the similar level of T-symmetry agreement explained by TDM as the data samples in  $\mathcal{D}$ .
Practical implementation. TSRL can be implemented based upon TD3 [2] by incorporating the proposed T-symmetry regularized representation and policy constraints as in Eq. (9) and (10), as well as the T-symmetry consistent latent space data augmentation. In our experiments reported in the next section, we generate  $K = 1$  augmented samples for each transition in the dataset, and filter based on the T-symmetry consistency loss. The pseudo-code of TSRL is summarized in Algorithm 1.
Algorithm 1 T-Symmetry Regularized Offline RL (TSRL)
Require: Offline dataset  $\mathcal{D}$ , encoder  $\phi$ , latent forward and reverse dynamics models  $f$  and  $g$  from TDM trained using objective Eq. (8).
1: Compute the T-symmetry consistency loss  $\ell_{T - sym}(\phi (s,a))$  (Eq. (7)) for all samples in  $\mathcal{D}$ , and set their  $\tau$ -quantile value as the augmentation threshold  $h$ .
2: Initialize the policy network  $\pi$ , critic networks  $Q$  and their target network.
3: for  $t = 1, \dots, M$  training steps do
4: Sample a mini-batch  $B$  of samples  $\{(s,a,r,s')\} \sim \mathcal{D}$  and compute their representations  $\{(z_s,z_a,z_{s'})\}$ .
5: //  $T$ -symmetry consistent latent space data augmentation
6: Generate  $K$  perturbed samples by adding perturbations  $\epsilon \sim N(0, 0.01\sigma_{z_s})$  on latent states  $z_s$  of each sample in  $B$ , where  $\sigma_{z_s}$  is the std of latent states in data.
7: Add augmented samples  $(z_{s} + \epsilon, z_{a}, z_{s^{\prime}} + \epsilon^{\prime})$  to  $B$  if satisfies  $\ell_{T\text{-sym}}(z_s + \epsilon, z_a) \leq h$ .
8: // Critic training with T-symmetry regularized representation
9: Update the value function  $Q$  based on the policy evaluation objective Eq. (9).
10: // Policy training with  $T$ -symmetry regularized policy constraints
11: Update the policy  $\pi$  based on the policy improvement objective Eq. (10).
12: Soft update the target networks.
13: end for
Table 1: Average normalized score on D4RL MuJoCo and Adroit tasks with full and reduced-size datasets. Some of the full dataset performance scores are reported from the IQL [15], MOPO [13], and DOGE [20] papers. Complete scores for Adroit-human and cloned tasks are included in Appendix C.
Task	Ratio	Size	BC	TD3+BC	MOPO	CQL	IQL	DOGE	TSRL(ours)
Hopper-m	1	1M	52.9	59.3	28.0	58.5	66.3	98.6 ± 2.1	86.7±8.7
1/100	10k	29.7±11.7	40.1±18.6	5.5±2.3	43.1±24.6	46.7±6.5	44.2 ± 10.2	62.0±3.7
Hopper-mr	1	400k	18.1	60.9	67.5	95.0	94.7	76.2±17.7	78.7±28.1
1/40	10k	12.1±5.3	7.3±6.1	6.8±0.3	2.3±1.9	13.4±3.1	17.9 ± 4.5	21.8±8.2
Hopper-me	1	2M	52.5	98.0	23.7	105.4	91.5	102.7±5.2	95.9±18.4
1/200	10k	27.8±10.7	17.8±7.9	5.8±5.8	29.9±4.5	34.3±8.7	50.5 ± 25.2	50.9±8.6
Hopper-e	1	1M	108.0	100.1	16.2±6.2	98.4	99.3	107.4 ± 3.6	110.0 ±3.3
1/100	10k	20.8±6.9	23.2±18.2	6.5±3.7	33.0±22.2	38.4±11.3	54.5±21.5	82.7±21.9
Halfcheetah-m	1	1M	42.6	48.3	42.3	44.0	47.4	45.3±0.6	48.2 ± 0.7
1/100	10k	26.4±7.3	16.4±10.2	-1.1±4.1	35.8±3.8	29.9±0.12	36.2±3.4	38.4±3.1
Halfcheetah-mr	1	200k	55.2	44.6	53.1	45.5	44.2	42.8 ± 0.6	42.2±3.5
1/20	10k	14.3±7.8	17.9±9.5	11.7±5.2	8.1±9.4	22.7±6.4	23.4±3.6	28.1±3.5
Halfcheetah-me	1	2M	55.2	90.7	63.3	91.6	86.7	78.7±8.4	92.0±1.6
1/200	10k	19.1±9.4	15.4±10.7	-1.1±1.4	26.5±10.8	10.5±8.8	26.7±6.6	39.9±21.1
Halfcheetah-e	1	1M	92.2	82.1	1.4±2.2	95.6	88.9±1.2	93.5 ± 4.2	94.3±5.5
1/100	10k	1.10±2.4	1.72±3.3	-0.6±1.1	4.2±0.94	-2.0±0.4	1.4±2.1	40.6±24.4
Walker2d-m	1	1M	75.3	83.7	17.8	72.5	78.3	86.8 ± 0.8	77.5 ±4.5
1/100	10k	15.8±14.1	7.4±13.1	3.1±4.7	18.8±18.8	22.5±3.8	45.1 ± 10.2	49.7±10.6
Walker2d-mr	1	300k	26.0	81.8	39.0	77.2	73.9	87.3 ± 2.3	66.1±12.0
1/30	10k	1.4±1.9	5.7±5.8	3.3±2.7	8.5±2.19	10.7±11.9	13.5±8.4	26.0±11.3
Walker2d-me	1	2M	107.5	110.1	44.6	108.8	109.6	110.4±1.5	109.8±3.12
1/200	10k	21.7±8.2	7.9±9.1	0.6±2.7	19.1±14.4	26.5±8.6	35.3 ± 11.6	46.4±17.4
Walker2d-e	1	1M	107.9	108.2	0.1±0.3	101.3	109.7±0.1	107.3±2.3	110.2±0.3
1/100	10k	10.4±5.3	23.8±16.0	1.4±3.4	41.6±21.6	12.6±4.5	72.1 ± 16.2	102.2±11.3
Adroit-human-total	1	5k	36.4	10.6	9.5	52.2	77.3	39.0 ± 17.1	80.9±21.1
Adroit-cloned-total	1	500k	57.5	41.1	-1.2	41.6	40.8	56.7 ± 16.2	66.3±23.7
1/50	10k	29.5±37.8	0.2±0.1	-1.7±1.5	0.6±0.8	32.7±24.6	29.7 ± 20.4	44.9±25.7
We evaluate TSRL on the D4RL MuJoCo-v2 and Adroit-v1 benchmark datasets [5] against behavior cloning (BC) as well as state-of-the-art (SOTA) offline RL methods, including model-free methods TD3+BC [8], CQL [9], IQL [15], DOGE [20], and model-based method MOPO [13]. In particular, we compare with DOGE [20], which leverages a state-conditioned distance function as policy constraint and exhibits strong OOD generalization performance. We report the final normalized performance of each algorithm after training 1M steps.
Performance on small datasets. We compare the performance of TSRL and the baseline methods on both the full D4RL datasets and their reduced-size datasets with only  $5\mathrm{k}\sim 10\mathrm{k}$  samples, which are constructed by randomly sampling a given fraction of trajectories in the full datasets*. These reduced-size datasets are only about  $1/20\sim 1/200$  of their original size. Compared with the performances on the full datasets, most baseline offline RL methods suffer from a noticeable performance drop under these extremely small datasets, mainly due to their over-reliance on the size and coverage of training data. Among all baselines, DOGE [20] is the only baseline that can still achieve reasonable performance in some tasks, primarily due to its capability to leverage the relationship between dataset geometry and the generalization capability of deep Q-functions, which is also beneficial for policy learning under small datasets. However, we still observe that DOGE struggles in some datasets, especially those with an extremely narrow distribution.
By contrast, TSRL achieves substantially better performance in all small dataset tasks, indicating superior sample efficiency. Moreover, although MOPO [13] also learns a dynamics model for offline policy learning, it performs badly when the dataset is small, revealing the importance of using a well-regularized model like TDM in the small-sample regime. It is interesting to see that most
*We didn't construct reduced-size datasets for Adroit-human tasks, as the full datasets are already very small.
Figure 3: Performance of TSRL and baselines algorithms on datasets with different training sample sizes
Figure 4: Impacts of different representation learning methods on 10k datasets.
methods fail on the 10k Halfcheetah-expert dataset, with only TSRL showing some success. We suspect this is due to the extremely narrow data distribution of the reduced dataset, which requires strong OOD generalization ability for algorithms to solve the task.
To further examine the impact of the training data size on algorithm performance, we also conduct experiments on three Walker2d datasets (medium, medium-expert, and expert) by varying the size of samples from 1M to 5k. The results are presented in Figure 3. It can be observed that most baseline offline RL algorithms experience a sharp performance drop when the datasets are reduced to 10k samples. Whereas, TSRL is still capable of preserving reasonable performance as the decrease of data size, even for extremely small datasets that contain only 5k samples.
Investigation on learned representations. To investigate the quality of the latent representation learned in TDM, we compare the performance of different representation learning approaches on the 10k datasets in Figure 4. To solely evaluate the impact of the representation, we remove the latent space data augmentation component from TSRL ("TSRL-no-A") and replace the state-action encoder  $\phi(s,a)$  learned from other representation learning approaches, including the autoencoder ("AE-rep"), autoencoder with latent forward dynamics ("AE-fwd-rep") without the ODE structure and the T-symmetry regularization in TDM, and a recent popular self-supervised representation learning method SimSiam [44] ("SimSiam"). To further investigate the impact of enforcing the ODE property, we also consider a variant of TDM ("TDM-no-ODE") by removing the ODE structure in latent forward and reverse dynamics. More detailed experiment setups are presented in Appendix B.
The results demonstrate that TDM representation achieves the best performance in all small-dataset experiments. By comparing "TSRL-no-A" and "TDM-no-ODE" with "AE-rep" and "AE-fwd-rep", we can see that the bi-directional design and the T-symmetry regularization are crucial for performance improvement. Moreover, we find "TSRL-no-A" consistently achieves better performance and lower variance as compared to "TDM-no-ODE", further confirming the benefit of incorporating ODE structure in producing a well-behaved representation for downstream tasks under small datasets.
Evaluation on data augmentation. We evaluate the impact of the proposed T-symmetry consistent latent space data augmentation in Figure 5. The results show that our proposed data augmentation scheme can help speed up convergence and reduce variance. Compared with less principled data augmentation methods such as adding zero-mean Gaussian noises as in S4RL [42], our method offers much better performance improvement. As shown in Figure 5, blindly adding random perturbations could suffer from performance degradation over the course of training, while TSRL with T-symmetry consistent data augmentation enjoys better training robustness. Additional comparative experiments on the proposed data augmentation method can be found in Table 6 of Appendix C.
Additional ablations. We also investigate the joint impact of the three design components in TSRL, including T-symmetry regularized representation and policy constraints, as well as the T-symmetry
Figure 5: Impact of T-symmetry consistent data augmentation on 10k datasets. "TSRL+S4RL" denotes replacing the latent space data augmentation in TSRL with the zero-mean Gaussian noise  $(N(0,0.1I))$ .
Figure 6: Ablation on TSRL on 10k datasets. "no-R": no T-symmetry regularized representation; "no-P": no T-symmetry policy constraints, and use BC term as in TD3+BC; "no-A": no latent space data augmentation.
Figure 7: Comparison of TSRL and baselines trained on the Walk2d-medium (top) and medium-expert (bottom) datasets that removing all samples with x-velocity of the top  $>0.2 \times$  max-x-velocity in the data. Left: learning curves. Right: x-velocity distribution of policy evaluation rollouts during the last 10k training steps.
consistent latent space data augmentation. We include TD3+BC for comparison, as it can be perceived as a vanilla version of TSRL without the previous three components. Moreover, the variant of TSRL with only the latent representation removed is not evaluated, as the latent space data augmentation also depends on the representation provided by TDM. Figure 6 presents the performance of all variants of TSRL on the 10k Halfcheetah medium and medium-expert datasets. As expected, it is observed that T-symmetry regularized representation and policy constraints jointly play a critical role in maintaining the performance of TSRL. As discussed previously, adding T-symmetry consistent latent space data augmentation also shows a positive impact on performance. Furthermore, it is observed that TD3+BC suffers from over-fitting on small datasets as its performance drops significantly with the increase of training steps, especially in the 10k Halfcheetah-medium dataset, while this phenomenon is not observed in TSRL. Nevertheless, the complete TSRL achieves the best performance in both tasks.
Generalization performance. To further verify the generalizability of TSRL, we construct low-speed datasets from the Walker2d-medium and Walker2d-medium-expert datasets by filtering out all high x-velocity samples (x-velocity of the top  $>0.2 \times$  max-x-velocity). This results in two smaller datasets with a large proportion of transition dynamics unobserved. We want to test if the agent can still generalize and learn well given only these low x-velocity data, with all high-speed samples removed. The experiment results are presented in Figure 7.
Under this setting, we observed that all existing offline RL baselines perform poorly when trained with only the low-speed dataset. This is primarily due to over-conservative data-related regularizations, which cause ineffective policy learning if the OOD region occupies the majority of state-action
space. CQL performs especially poorly in both tasks, perhaps due to over-conservative value function learning that impedes the policy to acquire some necessary control strategy to finish the task. TD3+BC performs poorly in the low-speed medium dataset, likely because this dataset has a narrower data distribution than the medium-expert dataset, and the latter could still contain some samples with reasonable speeds after filtering with x-velocity  $>0.2 \times \max -x$ -velocity in the dataset. IQL exhibits some level of generalization capability, but is still much weaker as compared to TSRL. By contrast, we observe that TSRL is still able to achieve good performance, due to the access to more fundamental dynamics information that remains invariant in both low- and high-speed data. This can be further verified if we inspect the policy rollout distributions (right figures of Figure 7) that the policies learned by TSRL indeed generalize to high-speed behavior that is not present in the training data.
Learning fundamental dynamics in physical systems. Learning conservation laws or invariant properties within a physical system is an active research area in physics [45, 46, 29, 30], climate science [47], and neuroscience [48], etc. A classic approach is based on Koopman theory, which represents the nonlinear dynamics in terms of an infinitedimensional linear operator [27]. In practice, this is achieved by finding a coordinate transformation to produce a finite-dimensional representation in which the non-linear dynamics are approximately linear. However, it also suffers from computationally expensive coordinate transformations and is only able to approximate the system dynamics. Another approach is utilizing a sparse regression model with the fewest terms to describe the nonlinear system dynamics [29, 30]. However, it assumes that the dynamical systems only have a few critical terms, which severely limits the model expressiveness and often requires prior knowledge of these critical terms. Based on expressive deep neural networks, a recently emerged research direction is to build ODE networks to learn conservation law in the dynamical system from data [49, 50, 51, 26]. Our proposed TDM falls within this direction, which models both forward and reverse latent ODE dynamics with deep neural networks and incorporates additional regularization on T-symmetry.
Offline reinforcement learning. Offline RL addresses the challenge of deriving policies from fixed, pre-collected datasets without interaction with the environment. Under this offline learning paradigm, conventional off-policy RL approaches are prone to substantial value overestimation when there is a large deviation between the policy and data distributions. Existing offline RL methods address this issue by following several directions, such as constraining the learned policy to be "close" to the behavior policy [6, 3, 8, 52], regularizing value function on OOD samples [9, 10, 11, 53], enforcing strict in-sample learning [16, 15, 17, 18, 54], and performing pessimistic policy learning with uncertainty-based reward or value penalties [13, 14, 4, 12, 55, 33]. Most existing offline RL methods adopt the pessimism principle and avoid policy evaluation on OOD samples. Although this treatment helps to alleviate exploitation error accumulation, it can be over-conservative and causes severe performance degradation if the training dataset is small or has poor state-action space coverage [20]. TSRL tackles this issue by allowing dynamics explainable OOD samples for policy optimization, thus offering greatly improved small-sample performance.
In this paper, we propose a physics-informed dynamics model TDM and a new offline RL algorithm TSRL, which exploit the fundamental symmetries in the system dynamics for sample-efficient offline policy learning. TDM embeds and enforces T-symmetry between a pair of latent forward and reverse ODE dynamics to learn fundamental dynamics patterns in data. The well-behaved representations and a new reliability measure for OOD samples based on T-symmetry from TDM can be readily used to construct the proposed TSRL algorithm, which achieves strong performance on small D4RL benchmark datasets and exhibits good generalization ability. There are also some limitations in our proposed approach. For example, in order to learn a well-behaved dynamics model, we introduced a set of dynamics and symmetry regularizations in TDM, which are beneficial to improve model generalization, but will lose some model expressiveness. However, we believe this can be a worthwhile trade-off between precision and generalization under small dataset settings, due to substantially improved model robustness.
This work is supported by National Key Research and Development Program of China under Grant (2022YFB2502904), and funding from Global Data Solutions Limited and Intel Corporation.
[1] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.
[2] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning, pages 1587-1596, 2018.
[3] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Systems, pages 11761-11771, 2019.
[4] Xianyuan Zhan, Haoran Xu, Yue Zhang, Xiangyu Zhu, Honglei Yin, and Yu Zheng. Deepthermal: Combustion optimization for thermal power generating units using offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 4680-4688, 2022.
[5] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
[6] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, pages 2052-2062. PMLR, 2019.
[7] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361, 2019.
[8] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[9] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2020.
[10] Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference on Machine Learning, pages 5774-5783. PMLR, 2021.
[11] Haoran Xu, Xianyuan Zhan, and Xiangyu Zhu. Constraints penalized q-learning for safe offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8753-8760, 2022.
[12] Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhi-Hong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. In International Conference on Learning Representations, 2021.
[13] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. In Neural Information Processing Systems (NeurIPS), 2020.
[14] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. In Neural Information Processing Systems (NeurIPS), 2020.
[15] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations, 2022.
[54] Xiangsen Wang, Haoran Xu, Yinan Zheng, and Xianyuan Zhan. Offline multi-agent reinforcement learning with implicit global-to-local value regularization. Advances in Neural Information Processing Systems, 2023.
[55] Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. Advances in neural information processing systems, 34:7436-7447, 2021.
[56] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1958, 2014.
Implementation details for TDM. As discussed in Section 3 and the illustrative Figure 2, TDM is modeled as a physics-informed reconstruction model with embedded ODE latent dynamics and T-symmetry preserving design. To ensure optimal training performance of TDM, we have included some additional implementation details below. More hyperparameters details of TDM are discussed in the next section.