Look Beneath the Surface: Exploiting Fundamental Symmetry for Sample-Efficient Offline RL
Authors: Peng Cheng
Xianyuan Zhan
Zhihao Wu
Wenjia Zhang
Youfang Lin
Shou cheng Song
Han Wang
Li Jiang
Venue: NeurIPS 2023 poster (2023)
Keywords: sample efficiency; offline reinforcement learning; fundamental symmetry
TL;DR: This paper provide a new insight that leveraging fundamental symmetry of system dynamics can substantially enhance offline RL performance under small datasets.

Abstract: 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

Look Beneath the Surface: Exploiting Fundamental Symmetry for Sample-Efficient Offline RL
Peng Cheng" 1.3, Xianyuan Zhan Zhihao Wut 1.3 Wenjia Zhang?,
Shoucheng Song Han Wang Youfang Lin Li Jiang"
1 Beijing Jiaotong University, Beijing, China Tsinghua University, Beijing, China
Beijing Key Laboratory of Traffic Data Analysis and Mining. Beijing, China Shanghai Artificial Intelligence Laboratory, Shanghai, China pcheng66126. com, chanxianyuaneee tsinghua, edu.a..edd.cmm zhwu, yflinbbjtu. edu. cn
Abstract
Offline reinforcement learning (RL) offers 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. worldw data collection is often expensive and uncontrollable, leading to small and narrowly cov- ered datasets and posing significant challenges for practical deployments offline insight that leveraging the fundamental sym-
RL. In this paper. provide metry of system dynamics can substantially enhance offline RL performance under small datasets. Specifically, propose 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 rep- resentations for small datasets and a new reliability measure for OOD samples based on compliance with the T-symmetry. These 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 extensive experiments, we find TSRL achieves great performance on small benchmark datasets with as few 1% of the original samples, which significantly outperforms the recent offline RL algorithms in terms of data efficiency and generalizability. Code is available at:
Introduction
The recently emerged offline reinforcement learning (RL) provides paradigm to learn policies from pre-collected datasets without the need to interact with the environments [[] 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 perfor- mance drop, as illustrated in Figure T 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, is often impractical or costly collect such a large amount of data, and the real datasets might only narrowly cover the state-action space. Clearly, learning reliable policies from
"qqual contribution.
CCoreesponiing Author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

small datasets with partial coverage has become one of the most pressing challenges for success- ful real-world deployments of offline RI.
walker 2d 2eddum---22
welker26-expert-22
10c
Normoltedd scores
Unfortunately, sample-efficient design consid- erations have been largely overlooked in the majority of offline RL studies. Pessimism is universally adopted in existing offline RL meth- ods and various forms of data-related regulariza- tions have been applied to combat the distribu- tional shift and exploitation error accumulation issues [, such as conservatively restricting policy deviation from the behavioral data [3]
20-
T00 DOL
K2L
80 105-5C CO.
19%
Figure 1: Performance of recent offline RI. methods and our proposed TSRL on the D4RL-MuloCo Walker2d medium and expert datasets when reducing the number of samples from 1M (full dataset) to 10k (1%)
6. 7 [), regularizing value function on out-
of-distribution (OOD) samples (9]
10
11
12
learning policy pessimistic MDP 4 or adopting strict sample learning [SS
13
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 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 Icarning 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?
16
In this paper, we provide a new insight that exploiting the fundamental symmetries in the system dy- namies 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 [1 22 It suggests that the underly-
ing laws of physics should not change under the rime-reversal transformation: - -t [22] 23]. Specifically, we 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, T-symmetry describes a funda- mental property of a system, hence learning it in principle would not require 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 detecting unreliable or nonegenerll state-action samples.
Based on these intuitions, we develop a physics-informed T-symmetry enforced Dynamics Model (TDM) to learn well-behaved and generalizable dynamics model with small datasets. TDM enforces the extended T-symmetry between pair of latent space forward and reverse dynamics sub-models, which are modeled 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 detect unreliable or non-generalizable samples, which can serve as 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-ane ahe-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 the first offline RL method that demonstrates promising performance on extremely small datasets.
Preliminaries
Offfine reinforcement learning. We consider the standard Markov decision process (MDP) set- ting (24], which is represented tuple M (S,A, P, p,)), where S and A are the state and

---

action spaces, r(s, a) is scalar reward function, P is the transition dynamics, is the initial state distribution, and ? E (0, 1) is discount factor. The objective of RL is to learn a policy x (als)) by maximizing the expected cumulative discounted return E,, E= 2'r(st, (z))], 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 = mig F. ((. BT (00(s,a))==((ss00)) ()))))))
(1)
where Q denotes a target Q-function, which is a delayed copy of the current Q-function; 13 is the Bellman operator, which is often used as the Bellman evaluation operator 1" B"*((s,) 0) = r(s, ylEa' Q (s', a') in many RL algorithms.
Under the offline RL setting, we are provided with a fixed dataset D = {50, ao: To, $1, ) J 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 о. To avoid this issue, widely used offline RL framework adopts the following behavior regularization scheme which regularizes the divergence between the learned policy ı and the behavior policy n of the dataset D:
===
π arg max EswD.ant((1)) [Q(s. a) D (a(: [)][ [s)]]
where D(-II-) is some divergence measures, which can have either an explicit B g 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 dynamical systems. Most real- wollworld dynamical systems with state measurement x E Ω on some phase space Ω can be modeled or approximated by the system of non-linear first-order ordinary differential equations (ODEs) as dx az Fx). where F' is some general non-linear, at least C' -differentiable valure function. First-order ODE systems are said to be time-reversal symmetric if there an invertible transformation r Ω + Ω2 that reverses the direction of time [22 26]: (x) - ([)()). If we define time evolution operator Us! Ω - Ω as Us:: x() H (x)(x(t)) = x(1 + At). Then T-symmetry implies that I o U, = U-, o r. 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 x = (s, a). We can slightly abuse the notations and denote s = de as the dinevverivetive of the current state s, which can be approximated the difference between the next and current states, i.e., 8 = 5 - 8. For a dynamical system that satisfies T-symmetry, it suggests that if we learn a forward dynamics F(s, a) = and a reverse dynamics G(s', a') = - & as a pair of first-order ODEs, we should have [ss, a) = G(s''a").
dicision-makinn perspective. it is known that T-symmetry broken by irreversible actions or some special dynamic processes (c.g.. frictional force against motion). Hence in this paper, consider more generic treatment by leveraging an alternative ODE reverse dynamics model G(s', a. = -& to establish the T-symmetry with the forward dynamics, ie., enforcing F(s, a) GG(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 construct well-behaved data-driven dynamics model and robust offline RL algorithm under small datasets.
However, from
sometimes be
3 T-Symmetry Enforced Dynamics Model
In this section, we present the detailed design of TDM, which is capable of learning a more fun- damental 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
3

---

V,z, 8
Encoder
Decoders
Reconstruction five ||s-- Vule, Oili l|a WGG2
(2 0) 8
decoding: fax
[ (.T(( D];
2 Latent dynamics
1/2
Forward ODE:
1x R60 |00
+
Latent reverse dynamics
Reveria ODE:
4ppa (al6aa) 06066.00
Va
Via (ta) -
Teymmetry:
lola, + &f f06
T-symmetry
Lyow
Σ,
/Rree tay Ryoe Eiy Eprype - A
V2IZA' (-5)
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 regu- larization, 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 mateeactioon encoder ((s, a) = (Z,, and pair of decoders (., 5,), th (zn) = a that embed the state-action pair (s, a) into latent representations (Zs, za) and then map them back. Specifically, we require the state decoder (-. 6,) be capable of decoding both 2u and 2.. where 5, is an indicator to help the decoder to decide the target output, with
5g as decoding 2s - and 8, as decoding & + s. The encoder & and decoders vh and Voe
induce the following reconstruction loss term for each state-action pair a):
(se a) = ||s (2,, 1) + ||a via (2a 112
Latent forward dynamics. Inspired by the prior works that incorporate physics- informed informa- tion into dynamical systems modeling [9] we embed a discrete-time first-order ODE system to capture the latent forward dynamics [( za) = á. Similar to é, we write é = 2x" 2s to denote the forward difference of the next and current latent state representations. Note that based on the chain-rule, we have é = de. Us.. de de
V. á. To enforce the ODE property, we can introduce
dr
8s
the following loss term for f:
(fwa(s, a, s') = (I 22,)8 2. 12 dō(s.a) f(ő(s, a))l
as
Minimizing Llud ensures that the latent forward dynamics correctly predicts the forward time evolution of latent states in the dynamical system. We also require the decoder (. 6%) to be able to decode é from é, to ensure it is compatible with the ODE property, which implies the following loss:
Las (8. a, s') å v,, (2., 1)112 = ||s - v. (f((()))))) 1)11
(5)
Latent reverse dynamics. We can further introduce a latent reverse dynamics g(zs/, 2a) - in the model, which captures the reverse time evolution of the system in the latent space. Similar to the forward dynamics loss Ldd we can write the reverse dynamics loss for g as:
06(s', a)
(rus a, s') I(V v 22 (- (-2) 2
(-8) g((((' ((
(6)
as'
T-symmetry regularization. The above latent forward and reverse dynamics 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,, =a)
-[ 2a). To further couple the learning process of f and g. note that Egg = 2, + 2. = 2, + f(z., 2n), which suggests g(2,', =,) = g(2. + [(ā,, Ra); za) - [[ 2.). This implies the following T-symmetry consistency loss:
&T- za) (z., 2n) g(2 + f(2., za),
2n)1 s)))))

---

Above instance-wist wise T-symmetry consistency loss also provides an alternative measure for evaluating the reliability of data sample. state-action pair (s, a) with large sy. (Ó(s, (o(ssoo a)) implies that this sample may not be well-explained by TDM or consistent with the fundamental symmetry of the system. This can used to detect unreliable OOD samples in offline policy optimization as well as construct 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:
LTDM
Σ 1Eree + las + fodd + lres + (ss s') (t) + (g)
(8)
(5,0
where (f) and (g) are Ll-noorms of the parameters of f and g. and λL1 is a scale parameter. L1 regularization is introduced to encourage learning parsimonious latent dynamies 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 13 14 31
35). The dynamics models in MBRL focus on constructing predictive model represent the forward transition dynamics of the system. Whereas, TDM 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 result, TDM can be substantially more well-behaved and robust when learning from small datasets.
32 33 34
4 T-Symmetry Regularized Offline RL
In this section, we discuss how incorporate the properties of TDM to construct 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 TDM also provides new forms of policy constraints to penalize unreliable and non non-generalizable OOD samples. Finally, compliance with T-symmetry enables a new latent space data augmentation procedure to further enhance the algorithm's 's performance with limited data.
T-symmetry regularized representation. fective approach to enhancing sample efficiency and generalization in many online and offline RI. studies [36] 37 38
Representation learning has been shown to be ef-
39 40]. A notable property of TDM that the learned latent state-action represen- tations from the encoder (Za, =a) = 6(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., zu) extracted by the encoder d(s, a) of TDM in the value function learning, which gives the following policy evaluation objective:
= argmin E(e,a,)')~ (r(*,a) a) + 70( -+1=-))) - ((o(s,))))
Q
T-symmetry regularized policy constraints. Existing offline RL methods primarily penalize the divergence between the learned policy and the behavioral data in the original action space, which ignores the underlying manifold structure of actions in the latent space [4]] and the system dynamics properties. Moreover, restricting policy optimization only within data-covered regions can be over- conservative 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:
argmax (s,a) (aQ(o(s, TT- Aillza- 2oll3 X2 (0(s, ((1s,
(10)
where latent actions 2, and z are obtained from o(s, a) and ø(s, (--ss) respectively. The second term restricts the latent action =,* of policy π from deviating too much from the latent action 2, in
data. The third term regularizes the T-symmetry consistency of policy-induced samples (s, n (:1s)). which is evaluated based on Eq. 7 and the learned TDM. A₁ and A2 are weight parameters, which only need to be roughly adjusted to ensure both the regularization terms are in similar scale
5

---

as the first term. We also introduce a normalization term a on the value function for training stability similar TD3+BC 8 which is computed based on a training batch B of samples as oo/ E(S.))BB Q(o(s, (1)(),(((1))))))))))) (-|+))]]] We set og = 2.5 in all of our experiments without tuning.
Instead of strictly regularizing OOD actions as in existing offline RL algorithms, our policy regu- larization 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 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.
T-symeerry consistent latent space data augmentation. It has been shown in previous studies 42 28 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 states 22 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 to a latent state Zg. i.c., (Ze, ka) + (z, e, =a), then the corresponding perturbation e' on the next latent state 2g' according to the latent forward dynamics & f(z.. =a) satisfies: ze' 2g + fz + 2,,) 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 - å. = g(Ee', za) as: 2, tc" = tc (g( z0).
Clearly, we should have c", which suggests the following condition:
e" - f(=, + 2,) + g(2. + f(z, + t, za), 2a) =0
(11)
This is exactly equivalent to requiring the instance-wist T-symmetry consistency loss (Eq. &y- (2, + c, =a) = 0. Hence we can use T-symmetry consistency loss (y-s as a reliability measure to filter out unreliable augmented samples (2. + C, zn) that inconsistent with the T- symmetry property of the learned latent dynamics in TDM. In our implementation, we only keep augmented samples that satisfy (T-sym (2, +e, S h. where we consider non-parametric treatment for threshold h, by setting it as the r-quantile value of all &T-ymm a)) values of (s, a) in D (we choose = 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 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 (00), as well as the T-symmetry consistent latent space data augmentation. In 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 I
Algorithm T-Symmetry Regularized Offline RL (TSRL)
Require:
Offline dataset D. encoder o, latent forward and reverse dynamics models f and g from TDM trained
using objective Eq. (B)
1: Compute the T-symmetry consistency loss Ey-sym (Φ((, a)) (Eq. for all samples in D, and set their T-quantile value as the augmentation threshold h.
2: Initialize the policy network म, critic networks Q and their target network.
3: (or(= 1,... .M training steps do
4: Sample a mini-batch B of samples ((s,a,r,s')) ~ D and compute their representations ((ze, Ray Zgg 2.))
5: //T-symmetry consistent latent spuce data augmentation
6: Generate K perturbed samples by adding perturbations
0.0l0,, latent states 2, of each
sample B. where d. is the std latent states in data.
7:
Add augmented samples (z. + E, ža, Z,' + é') to B if satisties fy (3+ + c, Ra) S h.
8:
// Critic training with T-symmetry regularized representation
9:
Update the value function Q based on the policy evaluation objective Eq.
9
10:
1/ Policy training with T-symmetry regularized policy constraints
11:
Update the policy æ based on the policy improvement objective Eq. (T0
12:
Soft update the target networks.
13:
end for

---

Table 1: Average normalized score D4RL MuJoCo and Adroit tasks with full and reduced-size datasets.
Some of the full dataset performance scores are reported from the IQL [5]. MOPO 13 and DOGE [20] papers. Complete scores for Adroit-human and cloned tasks are included Appendix
Task
Ratio
Size
BC
ID3+BC
MOPO
CQL
IQL
DOGE
(ORL(
1M
52.9
59.3
28.0
58.5
66.3
98.6 2.1
86,7+8.77
Hopper-m
1/100
10k
29.74:117
40.11:88.6
5.5
43.14:24.6
46.7::6.5
44.2 + 10.2
62.0::3.7
1
400k
18.1
60.9
67.5
95.0
94.7
76.2.:.7.7
78.7:1.8..1
Hopper-mr
1/40
5..1
7.346.1
6.8:00.3
2.3::199
13.4::311
17.9
21.8:88.2
Hopper-me
1
2M
52.5
98.0
23.7
105.4
91.5
102,7= 5.2
95 18.4
1/200
10.8
17.8=799
5.8::5.8
29.9--4.5
34.3:8.7
50.9:18.6
Hopper-e
1M
108.0
100.1
16.21:6.2
98.4
99.3
107.4 d: 3.6
110.0 3.3.3
1/100
10k
2..8
28.2
6.5.3..7
33.0:122.2
11.4111.3
54.51:21.5
21.9
42.6
44.0
45.3:: 0.6
0.7.
Hallcheetuh m
1/100
10k
7..4
16.4
-1.14.11
35.8.13.88
29.9:10.12
36.24:3.4
38.413.11
200k
55.2
44.6
53.1
45.5
44,2
42.8 =066
42.24:3.5
Halfchectah-mr
1/20
10k
14.3=7.8
17,9=9,5
11.75:2
8.14994
22.7+66.4
23.423.6
28.1+:3.5
Halfcheetah--me
2M
55.2
90.7
63.3
91.6
86.7
78.7+8.4
92.0-11.6
1/200
10k
19.1
19.4
15.
11.11:44
26.5
10.5.18.8
26.7
29.9
Halfcheerah-c
1
1M
92.2
82.1
1.4::2.2
95.6
88.9:11.2
93.5 4: 4.2
94.3-1:5,5
1/100
10k
1.10=2,4
1,72=3.3
-0.6+1.1
4.240.94
-2.0=0.4
1.4:2.1
40.6::24.4
Walker2d-m
IM
75.3
83.7
17.8
72.5
78.3
86.8 0.8
77.5 14455
1/100
10k
15.8
1.4
3.144.77
18.8
22.5+
45.1 # 10.2
10.7
Walker2d-mr
1
200k
26.0
81.8
39.0
77.2
73.9
87.3 2.3
66.1+112.0
1/30
10%
1,1199
5.715588
3.31:2.7
8.5:1
10:71:1199
13.5 8.4
26.0:1:1..3
Walker2d-me
107.5
44.6
108.8
109.6
110.4:1.55
109.811..1
1/200
10k
21.7=8.2
7.949.1
0.6.22.7
14.1
26.54:8.6
35.3 # 11.6
46.4
Walker2d-e
IM
107.9
108.2
0.1110.3
101.3
109.7:10.1
107.3:1..3
110.2.:033
1/100
10k
10.4=3.3
23.81:6.00
1.4:3.44
41.6::21.6
12.6:::455
72.1 116.
102.2::1133
Adroit-human-tar
1
5k
36.4
10.6
9.5
52.2
77.3
39.0 +: 17.1
80.9:121.1
Adroit cemit-cloned-tole
1
500k
57.5
41.1
-1.2
41.6
40.8
$6.7 1: 16.2
23.7
1/50
10k
29.5
0.240.11
-1.722.5
0.6::0.8
32.7:124.6
29.7 +: 20.4
44.9:1.5.7
5 Experiments
We evaluate TSRL the D4RL MuJoCo-v2 and Adroit-vl benchmark datasets [5] against behavior cloning (BC) as well as of-t--of-theeart (SOTA) offline RL methods, including model-free methods TD3+BC 8 CQL [9. IQL [15). DOGE (0). 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 IM steps.
Performance 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 5k~10k samples, which are constructed by randomly sampling given fraction of trajectories in the full dataset?"] These reduced- size datasets are only about 1/20~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 (2) is the only baseline that 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 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 mamlle regime. It is interesting to see that most
"We didn't 't construct seduced-size datasets for Adroit- human tasks, the full datasets already very small.
7

---

Normalized scores
walker2d-meeluum-2
Normalized scores
walkerr2d-meddumm-eeerttv
Normalized scores
walker2d-expert-22
BC 103
1BCMOTO
COL
ICL
DOGE
15AL
103 IBCMPFO
CD.
123.
DOGE
TS
BC
103 BLMOPO
CO.
12L.
DOGE
19AL
10k
Figure 3: Performance of TSRL and baselines algorithms datasets with different training sample sizes
Normalized scores
20-
halfcheetah-mediune
halfcheetah- hatheethh-mediuu--xeeerraa
walk 2ellkrr2--m
walkerrd mellkerdd-meeeer
Atrrep
Sm33am
Pi-lodrrep
10M-no-O0E
1SRL-nov/
Figure 4: Impacts of different representation learning methods on 10k datasets.
methods fail on the 10k Halfcheetah-exper 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 solve the task.
To further examine the impact of the training data size on algorithm performance, we also conduct experiments three Walker2d datasets (medium, medium-expert, and expert) by varying the size of samples from IM to 5k. The results are presented in Figure 3 It can be observed that most baseline offline RL algorithms experience sharp performance drop when the datasets are reduced 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, compare the performance of different representation learning approaches 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 ő(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--o-ODE" by removing the ODE structure in latent forward and reverse dynamics. More detailed experiment setups are presented 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- "AE-fwd-repp, we can see that the bi-directional design and the T-symmetry regularization are crucial for performance improvement. Moreover, find "TSRL-no-A" consistently achieves better performance and lower variance as compared to "IDM-no-ODE" further confirming the benefit incorporating ODE structure in producing 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 5 The results show that proposed data augmentation scheme can help speed up convergence and reduce variance. Compared with less principled data augmentation methods such adding zero-mean Gaussian noises in S4RL 42], our method offers much better performance improvement, As shown in Figure 3 blindly adding random perturbations could suffer from performance degradation consistent data augmentation enjoys better training robustness. Additional comparative experiments on the proposed data augmentation method can be found in Table Appendix C
the course of training, while TSRL with T-symmetry
Additional ablations. We also investigate the joint impact three design components in TSRL, including T-symmetry regularized representation and policy constraints, as well as the T-symmetry

---

Saaffcheeta-mmdiiunme
haffcheetah-mediuee expert-22
Taltcheetah-mmddumuy
halfcheerah wwd.imedddmmmm expert-y2
Normalieed Return
Normalized Deturn
Normaliced Deturnn
20-
Normalizet aedd
20
20
5
20-
8.00
Training Steps (x10%)
0.25
0.75
1000
8.00
Training Steps (x10%)
0.25
0.75
800
Training Steps (x109)
0.25
0.50
0.75
8.0
Training Steps (x109)
0.50
0.75
1.00
TSRL T5RL-S4RL
TSAL-n0-A
TSAL
T05
T3 AGM
150L 92
15M.
Figure 5: Impact of T-symmetry consistent data aug-
Figure 6: Ablation on TSRL on 10k datasets. "no-R": no T-symmetry regularized representation; "no-P": no T-symmetry policy constraints, and BC term in TD3+BC; "no-A": no latent space data augmentation.
mentation 10k datasets."TTRLL "TSRL+S4RL"" denotes re- placing the latent space data augmentation in TSRL with the aero-mean Gaussian noise (N 0.12)).
Walkerr2dmmedium>2
Distribution of x-velocity in policy rollouts
Normalized Reeurn 30 Rt
Probability Denaity
Training data
0.2 maxxxvelochy
8.30
0.25
0.50
0.75
1.00
5.0
0.0
2.5
Walker2d Walkeerzd-medlumm-exprrr-
Distribution x-velocity policy rollouts
Normalized Return
Probability Denelty
Training data
8
0.2 man-o-ellccky
it. 8.
0-
0.4
8.00
0.25
0.50
0.75
3.00
00
-25
0.
75
Training Steps (x10%)
X-velocity of the top
TD3+BC
COL
1OL
TSRL
Figure 7: Comparison of TSRL and baselines trained on the Walk22-medium (top) and medium-expert (hottom) datasets that removing all samples with x-vooocity of the top> 0.2 x max-x-veed 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 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 0 presents the performance of all variants of TSRL on the 10k Halfcheetah medium and medium-exper datasets. As expected, it is observed that T-symmetry regularized representation and policy constraints jointly play critical role in maintaining the performance of TSRL. As discussed previously, adding T-symmetry consistent latent space data augmentation also shows positive impact 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 observed in TSRL. Nevertheless, the complete TSRI. achieves the best performance both tasks.
Generalization performance. To further verify the generalizability of TSRL. construct low- speed datasets from the Walker2d-medium and fiikkerd-mediuumccne high x-velocity samples (x-velocity of the top>> 0.2x max- max-x-velocity). This results in two smaller datasets with a large proportion of transition dynamics unobserved. We want to test if the agent 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 /
datasets by filtering out all
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 wver-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.> 6.2xmax-x-veleee
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 ( that the policies learned by TSRL indeed generalizes to high-speed behavior that is not present in the training data.
Related Work
Learning fundamental dynamics in physical systems. Learning conservation laws or invariant
climate science [47,, and neuroscience [48]. etc. A classic approach is based on Koopman theory, which
properties within physical system is an active research area in physics 45
46
29 30
represents the nonlinear dynamics in terms of infinitedimensional linear operator 27 In practice,
this is achieved by finding a coordinate transformation to produce a miite-dimensionll representation in which the non-linear dynamics are approximately linear. However, it also suffers from computa- tionally expensive coordinate transformations and is only able to approximate the system dynamics. Another approach utilizing a sparse regression model with the fewest terms to describe the nonlinear system dynamics (29 BQ However, it assumes that the dynamical systems only have 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 RI. 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 52], regularizing value function on OOD samples [ 10 [3].
enforcing strict in-sample learning [6 15
17
18
(54) and performing pessimistic policy learning
55 [ 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 cver-conserv and causes severe performance degradation if the training dataset is small or has poor state-action space coverage 220.. TSRL tackles this issue by allowing dynamics explainable OOD samples for policy optimization, thus offering greatly improved small-sample performance.
with uncertainty-bat based reward or value penalties [1
14 12
7 Discussion and Conclusion
In this paper, we propose a physics-informed dynamics model TDM and 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 pair of latent forward and reverse ODE dynamics to learn fundamental dynamics patterns in data. The well- well-behaved representations and 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 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.
TDM, which are beneficial to
10

---

Acknowledgments and Disclosure of Funding
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.
References
[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. Advances Neural Information Processing Systems, pages 11761-11771, 2019.
[4] Xianyuan Zhan, Haoran Xu, Yue Zhang, Xiangyu Zhu, Honglei Yin, and Yu Zheng. Deepther- mal: 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, Ofrr 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-pooicy 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.11300 2019.
[8] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach 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 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-8766 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- 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 (NeurlPS), 2020.
[15] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit g-laaming. In International Conference on Learning Representations, 2022.
11

---

[16] David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline r1 without off-policy evaluation. Advances in Neural Information Processing Systems, 34:4933 4946, 2021.
[17] Haoran Xu, Li Jiang, Jianxiong Li, Zhuoran Yang, Zhaoran Wang, Victor Wai Kin Chan, and Xianyuan Zhan. Offline rl with no ood actions: In-sample learning via implicit value regularization. In The Eleventh International Conference on Learning Representations, 2023.
[18] Haoran Xu, Li Jiang, Li Jianxiong, and Xianyuan Zhan. A policy-guided imitation approach for offline reinforcement learning. Advances in Neural Information Processing Systems, 35:4085- 4098, 2022.
[19] Jinglin Chen and Nan Jiang. Information-theoreii considerations in batch reinforcement learning. In International Conference on Machine Learning, pages 1042-1051. PMLR, 2019.
[20] Jianxiong Li, Xianyuan Zhan, Haoran Xu. Xiangyu Zhu, Jingjing Liu, and Ya- Ya-Qin Zhang When data geometry meets deep function: Generalizing offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2022.
[21] JP Elliou and PG Dawber. Symmetry in Physics, Volume and 2. Macmillan, 1979.
[22] Jeroen SW Lamb and John AG Roberts. Time-reversal symmetry in dynamical systems: a survey, Physica D: Nonlinear Phenomena, 112(1-2):1-39, 1998.
[23] George Bluman and Sukeyuki Kumei. Symmetries and differential equations, volume 81.
Springer Science & Business Media, 2013.
[24] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[25] Ziyu Wang. Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems, 33:7768-7778. 2020.
[26] In Huh, Eunho Yang, Sung Ju Hwang, and Jinwoo Shin. Time-reversal symmetric ode network. Advances in Neural Information Processing Systems, 33:19016- 2020.
[27] Igor Mezić. Spectral properties of dynamical systems, model reduction and decompositions. Nonlinear Dynamics, 41:309-325, 2005.
[28] Matthias Weissenbacher, Samarth Sinha, Animesh Garg, and Kawahara Yoshinobu. Koopman q-learning: Offline reinforeement learning via symmetries of dynamics. International Conference Machine Learning, pages 23645-23667. PMLR, 2022.
[29] Steven Brunton, Joshua L Proctor, and Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences, 113(15):3932-3933 2016.
[30] Kathleen Champion, Bethany Lusch, J Nathan Kutz, and Steven L Brunton. Data-driven discovery of coordinates and governing equations. Proceedings of the National Academy of Sciences, 1 (16(45)::2444-224511 2019.
[31] Jianhao Wang, Wenzhe Li, Haozhe Jiang, Guangxiang Zhu, Siyuan Li, and Chongjie Zhang. Offline reinforcement learning with reverse model-based imagination. Advances in Neural Information Processing Systems, 34:29420-29432, 2021.
[32] Michael Janner, Justin Fu. Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, pages 12519-12530, 2019.
[33] Xianyuan Zhan, Xiangyu Zhu, and Haoran Xu. Model-based offline planning with trajec- tory pruning. In Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI-22, pages 3716-3722, 7 2022.
12

---

[34] Tao Yu, Cuiling Lan, Wenjun Zeng, Mingxiao Feng, Zhizheng Zhang, and Zhibo Chen. Playvir- tual: Augmenting sycle-consister virtual trajectories for reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[35] Changmin Yu, Dong Li, Jianye Hao, Jun Wang, and Neil Burgess. Learning state representations via retracing reinforcement learning. arXiv preprint arXiv:2111.12600, 2021.
[36] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine Learning invariant representations for reinforcement learning without reconstruction. 2021.
[37] Michael Laskin, Aravind Srinivas, and Picter Abbecl. Curl: Contrastive unsupervised represen- tations for reinforcement learning. In International Conference on Machine Learning, pages 1639-5650. PMLR, 2020.
[38] Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Marc G Bellemare. Contrastive hehavioral similarity emheddings for generalization in reinforcement learning. In International Conference Learning Representations, 2021.
[39] Mengjiao Yang and Ofir Nachum. Representation matters: offline pretraining for sequential decision making. In International Conference on Machine Learning, pages 11784-11794.
PMLR, 2021.
[40] Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. In International Conference on Learning Representations, 2022.
[41] Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforcement learning. In Conference on Robot Learning, pages 1719-1735. PMLR, 2021.
[42] Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple self-supervision for offline reinforcement learning in robotics. In Conference on Robot Learning, pages 907-917. PMLR, 2022.
[43] Jiafei Lyu, Xiu Li, and Zongqing Lu. Double check your state before trusting Confidence- aware bidirectional offline model-based imagination. Advances in Neurul Information Process- ing Systems, 35:38218 1--3211, 2022.
[44] Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 15750-15758, 2021.
[45] John David Anderson and Wendt. Computational fluid dynamics, volume 206. Springer, 1995.
[46] Ilia Grigorenko and Elena Grigorenko. Chaotic dynamics of the fractional lorenz system. Physical review letters, (3))003411 2003.
[47] Kevin E Trenberth and Kevin E Trenberth. Climate system modeling. Cambridge University Press, 1992.
[48] Eugene M Izhikevich. Dynamical systems in neuroscience. MIT press, 2007.
[49] Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
[50] Emilien Dupont, Arnau Doucet, and Yee Whye Teh. Augmented neural odes. Advances in Neural Information Processing Systems, 32, 2019.
[51] Xuanging Liu, Tesi Xiao, Si Si, Qin Cao, Sanjiv Kumar, and Cho-Jui Hsieh. Neural sde: Stabilizing neural ode networks with stochastic noise. arXiv preprint arXiv:1906.02355. 2019.
[52] Haoran Xu, Xianyuan Zhan, Jianxiong Li, and Honglei Yin. Offline reinforcement learning with soft behavior regularization. arXiv preprint 1rXiv:2/10.07333 2021.
[53] Haoyi NIu, Yiwen Qiu, Ming Li, Guyue Zhou, Jianming HU, Xianyuan Zhan, et al. When to trust your simulator: Dynamics-aware hybrid offline-and-online reinforcement learning. Advances in Neural Information Processing Systems, 35:36599-36611 2022.
13

---

[54] Xiangsen Wang, Haoran Xu, Yinan Zheng, and Xianyuan Zhan. Offline multi-agent reinforce- ment 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-basee 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, llya 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.
14

---

A Implementation Details
Implementation details for TDM. As discussed in Section and the illustrative Figure TDM is modeled as 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.
Network structure: In all our D4RL experiments, we implement the encoder, decoders, latent forward and reverse dynamics as 4-layer feed-forward neural networks with ReLU activation, and optimized using Adam optimizer. For the state decoder (5 6s), we concatenate an extra indicator 6, in the input to help the state decoder to decide the target output. More specifically, to decode 2s +5, we concatenate 6, = 0 with 2 as input; and for 8 + we concatenate 6, 1 with E.-
Computing second derivative of c(-): As TDM involves pair of latent ODE forward and reverse dynamics models, whose training losses Eq. 4 and (6) involve regressing on 3á(s.a) and Bélss', .a) (-)) as target values. This results in gradient through a gradient of (c)). Computationally, Os'
85
we calculate the Jacobian matrix Do(s.o) using the vmap() function in Functorch to ensure the second derivative of ((-) can be correctly backpropagated during stochastic gradient descent. Similar treatment can also be implemented with other diteediiferentiationn frameworks like Jax that support computing higher-order derivatives.
Pre- tre-training the encoder and decoders: As the final learning objective of TDM Eq. 8 involves several loss terms, we observe that in small datasets, loss terms such as the reconstruction loss (Eq, (3) for the encoder and decoders converges much slower than other loss terms. When updating all the loss terms with the same number of training steps, some losses suffer from over-fitting while others are still not fully converged. For these cases, we pre-train the encoder and decoders with the reconstruction loss for a given number of training steps, and then use the complete learning objective of TDM (Eq. (4) for the rest of the training. The numbers of treetririing/rriiingg epochs for the experiments in this paper are reported in Table
As reported in Table 2 we find that the number of pre-training epochs required for TDM to reach the best learning performance is associated with the specifie task and the size of training data. For small datasets, TDM generally needs more training and pre-training epochs to avoid overfitting the latent dynamics and T-symmetry losses. For MuJoCo locomotion tasks, recommend pre-training the encoder and decoders for 10% of the total training epoch. For the more complex adroit tasks, TDM requires more epochs to extract the ODE dynamics and T-symmetry property of the system dynamics. In this case, there is no pre-training necessary for the encoder and decoders.
Table 2: Training epochs of TDM for D4RL tasks with different dataset scales
Locomotion Tasks
Adroit Tasks
5k& 10k
50k 100k
Full dataset
5k&
Full dataset
Training epoch
2000
1000
200
2000
200
Pre-train epoch
200
100
20
0
0
Enhancement on the T-symmetry regularization: We observe that in some small datasets (mainly in the Halfcheetah environment), the training of the latent reverse dynamics model g might suffer from a certain level of degeneration. This is reflected the g(z, + [zs; Aa): 2,); =a) produces similar values as - (2,,, ),), resulting in small T-symmetry consistency loss values (Fq.. (7)). however, the discrepancy between ( =a) and [( za) remains large. To solve this issue and further enforce the T-symmetry, we apply the following enhanced T-symmetry regularization when such a phenomenon is observed:
LEnhanced-T sym (ž., za) = Ilf =a) + g(z. f(E,, za), 2a)1l + llf( =a) + ,( (12)
We find that applying the above enhanced T-symmetry loss can successfully resolve the degeneration issue of the latent reverse dynamics model and achieve good performance in the downstream offline
hatps///ithubb.om//ooggggg
15

---

RL tasks. However, we find in most small datasets, the original T-symmetry consistency loss is sufficient. We advise only to use the above enhanced T-symmetry consistency loss when large discrepancies between (zx: =a) + g(z. + f(zs: za), za) ll and [z za) + g za l are observed.
Hyperparameter details for TDM and TSRL. The architectural parameters of TDM and TSRL, as well as the TSRL. hyperparameters are summarized in Table Based on the different scales of the datasets, we basically only two of hyperparameters for all D4RL-MuJoCo locomotion tasks and only one set of hyperparameters for all D4RL-Adroit tasks. Because of the extremely narrow distribution of the reduced-size expert dataset, apply Dropout (56) with dropout rate of 0.1 to regularize the policy network in all tasks with 10k expert data.
Table 3: Hyperparameter details for TDM and TSRL
Hyperparameters
Value
Optimizer type
Adam
Weight of fr -ym and Eds and free
TDM Architecture
Weight of Eres and fy
0.1
Learning
x 10-4
State normalization
True
Hidden units of forward and reverse model
512
Hidden units of encoder
512 x 256 x 128
Critic neural network layer width
512
Actor neural network layer width
512
State normalization
True
Actor learning rate
x 10
Critic learning rate
3 10
TSRL Architecture
Policy noise
Policy noise clipping
Policy update frequency
2
Discount factor
0.99
Number iterations
10"
Target update rate
0.005
AL1
le-5
a
2.5
50% for Walker2d and Adroit tasks, 70% for HalfChectah and Hopper2d
TSRL
Hyperparameters
MuJoCo: 5 or 10 for full dataset, 100 or 200 for 10k dataset Adroit: 10,000 for both full and reduced datascts
A
for MuloCo full & Adroit datasets
100 for MuJoCo 10k dataset
B Detailed Experiment Setups
Reduced-size dataset generation. To create reasonable reduced-size D4RL datasets for fair comparison, we sub-sample the trajectories the datasets rather than directly sampling the (s, a, s',r) transitions. For example, there 2M (s, 8'r)) transitions in the maalfcheetah-medium-eeee dataset, we first split these records into 2,000 trajectories based on the done condition, then randomly draw 10 trajectories (10k transition points) serve the reduced-size datasets for model training.
Experiment setups for representation learning evaluation. To evaluate representation quality and the impact of each design choice of TDM, we compare TDM representation with several baselines on the small dataset settings. We provide the detailed description of the representation learning baselines follows:
16

---

"AE-rep" model: We construct a vanilla auto-entt eito-encoder without any further constraints during the learning process, which was trained by the reconstruction loss only. The network sizes of the encoder and decoders are the same as the ones used in TDM.
"AE-fwd-rep" model: Similar to the "AE-rep" model but with latent forward dynamics prediction model, f. which implemented as 4-layer forward neural network with ReLU activation, and optimized using Adam optimizer (same TDM). The forward model
trained by minimizing the loss term 2 - f(o(s, a))l where we directly regress f(o(s, a)) with the &g derived from the latent states obtained from the encoder as &, 2 Zg. Note that in this baseline, no ODE property nor T-symmetry regularization is included. Again we use the decoder to decode 2, + in TDM for the next state prediction.
"TDM-no-ODE" model: Holds the same structure with TDM but trained with no ODE property. More specifically, similar with "AE-fwd-rep", the latent forward and reverse dynamics model was trained by || [( a))l and (1(-).) g( a)) [β.. where -2, -, is directly calculated from the encoded latent states, i.e., 2, = Es¹ implicitly captured, since both the latent forward and reverse dynamics models are regressing the same & and its opposite value.
Note that in this baseline, the T-symmetry is also
For the self-super representation learning baseline, implement auto-encoder structure with the optimization objective proposed in the SimSiam paper 44] For detailed model description and hyperparameters setting, please refer to Chen et al. [4
"SimSiam" model:
Experiment setups for evaluating generalization performance. To evaluate TSRL's generaliza- tion capability beyond the offline datasets, we construct two low-speed datasets based on the original D4RL Walker2d medium and medium-exper datasets. In accordance with the Gym documentation, we selected the "x-coordinate velocity of the top" (8th dimension of the states) in the walker envi- romment perform data filtering. We remove all samples with the x-coordinate velocity of the top greater than 0.2x max-x x-eeoccyty recorded in the data. This results in two smaller low-speed datasets (about 200k for the medium dataset and 250k for the medium-expert dataset). We train TDM and TSRL on these low-speed datasets and the results are reported in Figure 7 (main paper).
C Additional Results
Complete results on D4RL Adroit tasks. The complete results of TSRL in Adroit human and cloned tasks with different dataset scales are presented Table I As shown in the results, TSRL achieves much better performance in the pen tasks, both the full datasets and the reduced-size datasets.
Table 4: Complete results on D4RL Adroit tasks
Task
Ratio
Size
BC
TD3+BC
MOPO
CQL
IQL
DOGE
TSRL
Pen-human
5k
34.4
8.4
9.7
37.5
71.5
42.6 + 16.3
80.1+18.1
Hammer-human
5k
1.5
2.0
0.2
4.4
1.4
-1.2 + 0.2
0.2 0.3
door-human
5k
0.5
0.5
-0.2
9.9
4.3
-1.1 4: 0.2
0.5 0.3
Relocate-human
5k
0.0
-0.3
-0.2
0.2
0.1
-0.3 + 0.5
0.1 + 0.1
500k
56.9
41.5
-0.1
39.2
37.3
56.9 1:: 15.2
64.9 20.1
Pen-cloned
I/50
10k
37.4 1. 37.6
-0.1 1 6.9
-0.1 + 0.1
1.5 i 4.8
35.6 1: 30.5
30.1 +:: 19.7
41.6 1. 27.5
500k
0.8
0.8
0.2
2.1
2.1
0.2 + 0.3
1.7 + 1.9
Hammer-cloned
1/50
10k
0.3 + 0.4
0.2 +: 0.1
0.1 + 0.1
0.2 + 0.1
0.4 1 0.2
0.3 0.1
0,6 +: 0.3
500k
-0.1
-0.4
-0.1
0.4
1.6
-0.1 0.1
-0.1 + 0.6
Door-cloned
1/50
10k
-0.1 + 0.1
-0.3 + 0.1
-0.2 0.1
-0.3 0.1
1.5 + 0.8
-0.5
-0.1 0.3
500k
-0.1
-0.3
-0.3
0.1
-0.2
-.22
-0.2 + 0.1
Relocate-clon clonee
1/50
10k
-0.2 + 0.1
-0.3 + 0.1
-0.3 0.1
-0.34 0.1
-0.1 + 0.5
-0.2
-0.2 0.1
Additional results on Antmaze-umaze tasks. umaze tasks with full and reduced-size 10k datasets. The results presented in Table We use the
We also conduct experiments on the D4RI Antmaze-
17

---

Table 5: Results D4RL Antmaze-umaze tasks with full and reduced-size datasets
Task
Ratio
Size
BC
TD3+BC
CQL
IQL
DOGE
(our)
54.6
78.6
84.8
85.5
97.0 1.8
81.4 + 19.2 15.1
Antmaze-u
1/100
10k
44.7 +42.1
0.7 + 1.2
0.1 0.0
65.1 = 19.4
56.3 + 24.4
1
1M
45.6
71.4
43.4
66.7
63.5 + 9.3
76.5 + 29.7
Antmaze-u-d
1/100
10k
24.1222.2
16.27 + 16.4
0.5 + 0.1
34.6 = 18.5
41.7 = 18.9
52.2 1:22.1
Table Performance data augmentation methods with 10k reduced- size D4RL datasets.
Task
CABI
S4RL-N
S4RL-U
TSRL
Hopper-m
48.3 + 3.9
28.3 + 6.2
23.6 = 4.7
62.0 + 3.7
Hopper-mr
19.8 1: 3.9
16.6 = 12.9
12.5 + 12.2
21.8 + 8.2
Hopper-me
38.3 + 5.3
12.5 + 3.8
13.1 = 4.7
50.9 8.6
Hopper-e
34.6 24.4
14.1 = 12.9
12.2 11.6
82.7 21.9
Halfcheetah-m
34.8 1.9
25.1 6.8
23.2 7.1
38.4 3.1
Halfcheetah-mr
23.5
15.1 + 9.3
14.8 9.5
28.1 3.5
Halfcheetah me
29.9 1.7
27.1 + 7.1
23.4 8.2
39.9 = 21.1
Halfchectah-e
4.2 122:44 4.1
2.4 = 3.9
1.8 = 3.1
40.6 + 24.4
Walker2d-m
42.4 #: 23.3
24.5 + 4.3
21.9 4.8
49.7 + 10.6
Walker2d-mr
11.7 7.6
1.5 = 2.1
1.4 + 2.3
26.0 = 11.3
Walker2d-me
17.4 + 9.2
21.9 = 16.4
16.0 + 13.2
46.7 + 17.4
Walker2d-e
20.2 + 3.4
56.5 = 26.7
51.1 + 29.7
102.2 = 11.3
same hyperparameters as in the D4RL MuloCo tasks. Again, we find that TSRL achieves comparable performance as other baselines on the full datasets, but is substantially better under small datasets.
Additional comparative results on data augmentation. We conduct additional experiments on the reduced-size 10k MuJoCo datasets to compare TSRL and other offline RL methods using data augmentation. In particular, we compared with model-free method S4RL 42 with Gaussian (S4RL-N) and uniform (S4RL-U) noises, as well as recent model-based data augmentation method CABI
4 CABI employs a pair of predictive dynamics models to assess the reliability of the augmented data. The results are presented in Table
The results clearly show that TSRL outperforms all offline RL baselines with data augmentation under small datasets. It also observed that model-based methods TSRL and CABI generally perform better than model-free data augmentation method S4RL in this setting, due to access to additional dynamics information. Moreover, as CABI does not learn a strongly regularized dynamics model with T-symmetry consistency as in our proposed TDM, it still has a noticeable performance gap as compared to our method.
Ablation the level of ODE and T-symmetry regularization in TDM. As discussed in the conclusion section of the main paper, TDM adds extra ODE dynamics and symmetry regularizations, which are beneficial to improve model generalization. but will lose some model expressiveness the regularization is too strong. In this section, we conduct an ablation study on the impact of the regularization strength of the ODE property and the satisfaction with the T-symmetry. Specifically, we vary the loss weights of L fres and f-- in the TDM learning objective (Eq. (8)) and train a loosely regularized and a strongly regularized TDM model on the 10k datasets (see Table (7) The loosely regularized model has the maximum reconstruction expressivity but may not produce a well-behaved representation due to weak regularization. Whereas the strongly regularized model sacrifices the expressivity for regularized behaviors. We further evaluate their performance with TSRL,, with the results reported in Table [ The experiment results demonstrated that an overly expressive model could not help the RL algorithm to derive well- belvebe policy with limited data
18

---

Table 7: TDM with different regularization strengths
Different versions TDM
Eyee
td
Erea
Ey-
AL1
Loosely regularized
0.01
0.01
0.01
le-5
Paper
0.1
0.1
0.1
1c-5
Strongly regularized
1
le-5
Table 8: Performance of TSRL with different TDM models on 10k datasets
Task
TDM (loosely regularized)
TDM (paper)
TDM (strongly regularized)
Hopper-m
13.7113.66
62.013.77
43.6.114.3
Hopper-m-r
15.4+9.7
2..8+
15.6+9.8
Hopper-m-e
49.7.117.1
50.9.1.66
30.9.120.5
Halfchectah-m
39.1+3.66
38.4#3.1
36.6++30.0
Halfcheetah-m-rr
28.31.6.9
28.11:3.5
22.9.18.4
ee
36.2+:5.4
39.9:t21.1
31.0 + 3.4
Walker2d-m
43.2:1:27.3
49.7.1:00.6
35.6:1:2622
Walker2d-m-r
20.2:4.1..1
26.0:t:1.3
21.7:6..1
Walker2d-m-e
25.9:1200.7
46.4:17.4
29.4:24.77
due to potential overfitting and inconsistency with the T-symmetry property. On the other hand, an overly regularized model may also hurt performance. This consistent with our previous insight that trade-off exists between model expressiveness and T-symmetry agreement. A proper balance between these two behaviors can be necessary for small-sample learning.
Learning curves for TSRL on D4RL locomotion tasks. The learning curves for reduced-size D4RL MuJoCo datasets with 10k samples are showed Figure 8 For each evaluation step, the policies are evaluated with 5 episodes over 3 random seeds.
19

---

Halfcheetah-expert
taaicchetta-mmedlu-exxxeeeerrr
Halfchhetth-mmeiuu--reepaaayy
Norma lzzd Return
Norma lizeed Return
Norma lized Return
8
80
10
-0
30
8.00
0.50
8,00
0.25
0.25
Training Steps (x10%)
Training Steps (x10%)
Training Steps (x10%)
Halicheetah-medium--2
Hxpper-expert-v2
Hopper-mediim--xxprrtt>>
Normalized Return
Normallzzen Retumm
Normalized Return
15
8 8 B 3 88
MAN Ma
30
10
8.20
8.0
Training Sceps (x10%)
Training Steps (x119))
Training (x109)
Hopper-medium-replavyy
8 a 8 to
Walker2d-experts22
Normalzed Retur
Nermalized Return
Normalized Return
50
vo-
30
00
6.00
0.25
0.50
0.75
1,00
8.00
0.25
8.50
0.25
3.00
8.00
0.50
0.75
1.00
Training Steps (x10%)
Training Steps (x10*)
Training Steps (x10*)
milkkrrd--ediimm-ee
Walker2d -eaieer--meduuum-eplaayy22 replay.v2
Walker2d-mediumyy
B 28 B to B 8
Normalized Retur
Rodda Return
lloraalized Return
5t 8 1
10
8.00
0.52
0.75
1.00
0.25
0.50
0.75
3,00
0.25
0.75
1.00
Training Steps (x10%)
Training Steps (x10%)
Training Steps (x10%)
Figure 8: learning curves for reduced-size D4RL MuJoCo datascts. Error bars indicate min and max values over 3 random seeds.
20


Reviews:
[{"review_id": "nDfby9s5eF", "reviewer": "Reviewer_P7X7", "summary": "The paper investigates the time-reversal symmetry of forward and reverse dynamics in reinforcement learning (RL). The authors propose a Time-reversal symmetry enforced Dynamics Model (TDM) that models the consistency between forward and reverse dynamics. Using the TDM, they further propose an offline RL algorithm that leverages the learned TDM in three ways: 1. Using the representation from TDM for value function learning; 2. Using TDM to penalize OOD samples; 3. Using TDM to moderate useful data augmentation. Extensive experiments demonstrate that the proposed method outperforms a number of baselines, especially, the proposed method can learn a better policy with significantly fewer samples.", "questions": "1. I wonder how the authors ensure the fairness of their experiments, given that the proposed method appears more complex than the baselines. Do all the baselines have a similar amount of trainable parameters as the proposed method?\n2. Instead of ensuring the T-symmetry using an extra loss term, is it possible to constrain the neural network architecture so that the model respects the T-symmetry by definition (like an equivariant neural network architecture)?\n3. The proposed method doesn't seem to be specific to offline RL. Have the authors tried to use it in online RL?", "limitations": "The authors do address the limitations of their work. However, the discussion could be more comprehensive. For example, under what circumstances would the T-symmetry break? Are there scenarios where the proposed method might fail?", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is novel and backed by convincing experimental results.\n2. The concept of time-reversal symmetry could potentially be widely applied in RL, as it represents a fundamental structure in many RL problems.\n", "weaknesses": "1. The proposed method leverages a learned dynamic model for offline RL, whereas most of the baselines are model-free (except for MOPO), which seems somewhat unfair. In particular, the proposed method appears somewhat related to the Dreamer [A, B] method, which also leverages a dynamic model for data augmentation. I am curious about the authors' thoughts on comparing with Dreamer.\n2. The claim of leveraging time-reversal symmetry is slightly concerning as the method essentially learns the forward and reverse dynamic models instead of leveraging the time-reversal symmetric property of the MDP (e.g., P(s, a, s')=P(s', a, s) in certain scenarios). \n\n[A] Hafner, Danijar, et al. \"Dream to Control: Learning Behaviors by Latent Imagination.\" International Conference on Learning Representations. 2019.  \n[B] Hafner, Danijar, et al. \"Mastering Atari with Discrete World Models.\" International Conference on Learning Representations. 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We really appreciate the reviewer for the positive feedback and valuable comments.\n\n> **W1. Comparison mostly to model-free methods rather than model-based methods. In particular, the comparison with Dreamer.**\n\n- First, we'd like to highlight that our approach is very different from the typical model-based methods. Most existing model-based methods use the learned dynamics model to generate imaginary rollouts to facilitate policy learning. Whereas in our approach, we **do not** use the model for rollouts generation, but rather only used it to learn well-behaved representations as well as consistency metrics for policy constraints and data augmentation. If the reviewer inspects our value and policy learning procedure in Section 4. it actually shares more similarities with many model-free policy constraint offline RL methods. \n- We adopt such a design as under offline small dataset settings, it is generally not possible to learn an accurate dynamics model for reliable rollout generation. A poor dynamics model can negatively impact policy learning. However, learning reasonable representations and consistency check metrics still remains possible if we learn a more fundamental dynamics model with strong physics-informed regularizations. This is evident as shown in Fig. 1, 3, and Table 1 in our paper, that model-based offline RL method like MOPO suffers from severe performance degradation with reduced data size as compared to model-free methods and our proposed TSRL. \n- Lastly, Dreamer v1&v2 are online RL methods, which allow continually acquiring fresh environment samples to improve the accuracy of the model. While in our setting, only a very small offline dataset is given. This significantly exacerbates the difficulty of model and policy learning. Note that even being sample-efficient, in the Dreamer paper, it still needs $5\\times 10^6$ online environment samples to learn good policies for many tasks, while in our setting, we only provide each algorithm 10k offline samples.\n\n\n> **W2. Not leveraging the time-reversal symmetric property of the MDP (e.g., P(s, a, s')=P(s', a, s) in certain scenarios)**\n\n- If the reviewer closely inspects our proposed extended T-symmetry, we introduce two ODE systems $F(s,a) = \\dot{s}$ and $G(s',a)=-\\dot{s}$, where $\\dot{s}=s'-s$. This is equivalent to constructing two consistent systems of $(s,a)\\rightarrow s'$ and $(s',a)\\rightarrow s$, essentially similar to the $P(s, a, s')=P(s', a, s)$ relationship in MDP mentioned by the reviewer, with the difference that we further require both dynamics are ODE systems.\n- Due to the above construction, our proposed extended T-symmetry is more broadly applicable as compared to the original T-symmetry (i.e., $F(s,a) = \\dot{s}=-G(s',a')$) in MDP settings, as the latter can sometimes be broken by irreversible actions. Please refer to L113-129 in our paper for a detailed discussion.\n\n> **Q1. Fairness regarding using more complex model. Do all the baselines have a similar amount of trainable parameters as the proposed method?**\n\n- In this paper, we propose an RL algorithm rather than a supervised learning model, the setting is different and the number of parameters does not necessarily determine the performance. Note in the Dreamer v1&v2 papers, these more complex methods are also compared with lightweight methods like DQN and A3C.\n- To ensure comparability, we use the same architectures for Q-networks and policy network as other baseline algorithms in our experiments. The only addition is the incorporation of TDM for representation learning and consistency metric computation. Note that simply adding extra model/parameters in representation learning may not yield comparable performance improvements as in TSRL. As reported in Fig. 4 of our paper, baselines using SimSiam and AE-fwd-rep representations have a similar amount of parameters as in TSRL, but their performance improvements are less significant.\n\n> **Q2. Possibility of constraining the neural network architecture to enforce T-symmetry by definition (e.g. equivariant NN)?**\n\n- We appreciate the reviewer for the insightful comment. We actually explored the possibility of employing equivariant NN architectures to enforce T-symmetry, but we found some applicability issues:\n    - Most existing equivariant NN techniques leverage relatively simple and explicitly known equivariant mapping/transformation of the system. However, in our problem, the forward and reverse ODE dynamics are unknown and need to be learned. Enforcing equivariant relationships between two unknown systems is less straightforward to be implemented in a NN architecture.\n- In contrast, our method embraces simplicity and generality, without relying on excessive data or task-specific knowledge. We enforce T-symmetry by simply incorporating a few supervised loss terms. Nevertheless, constraining the NN architecture is a valuable idea, we will continue to explore this direction in our future works.\n\n\n> **Q3. Applicability to online RL?**\n- In principle, our method can also be used in online RL. However, as discussed in W1, online RL allows continuous collection of fresh samples through online interaction, hence the accuracy of the model can be improved during training. This makes heavy regularization used in TDM less necessary for online settings, especially in the later part of the training.\n- We focus on the small-sample offline RL setting, as it poses some more demanding challenges, and has many real-world deployment scenarios. Only a very small number of samples are given, one has to incorporate strong regularization to ensure reasonable generalization as well as alleviate distributional shift in offline policy optimization.\n\n\n> **Limitations: more discussion on the cases when T-symmetry breaks**\n\n- Please refer to our response to W2. Our proposed extended T-symmetry improves over the original T-symmetry definition for generic MDP settings and is not susceptible to the negative impact of irreversible actions."}, {"review_id": "iFFTNj3w18", "reviewer": "Reviewer_5ovg", "summary": "This work introduced a Time-reversal symmetry enforced dynamics model, which leverages the consistency between a pair of forward and reverse latent dynamics for improving the sample efficiency of offline RL algorithms. Conducted experiments demonstrate the effectiveness of the proposed method.", "questions": "1. Why can the proposed method improve the sample efficiency of off-line RL algorithms? Pls give in-depth and convincing analysis.\n2. What are the relations between this work and those metioned in the weakness part?\n3.An open question: Are there any ideas to extend the core idea of this work to Online RL algorithms?", "limitations": "Pls see the weakness and questions parts.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method makes sense in improving the sample efficiency of offline RL algorithms.\n- This work is clearly presented so that it is easy for readers to catch up with the main ideas.", "weaknesses": "- I suggest authors to strengthen the analysis on the rationale, i.e., why can the proposed method improve the sample efficiency of off-line RL algorithms.  This will make the paper more insightful.\n- This work lacks sufficient discussion on the relation/comparison with related works, especially for those also exploiting the consistency, e.g.,    PlayVirtual [1] and some others cited by [1].\n[1] Yu, Tao, et al. \"Playvirtual: Augmenting cycle-consistent virtual trajectories for reinforcement learning.\" Advances in Neural Information Processing Systems 34 (2021): 5276-5289.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **W1 & Q1. Rational of the sample efficiency improvement of TSRL**\n\nThe sample efficiency of TSRL is a joint result of a series of elegant and closely related design choices:\n\n- Firstly, learning fundamental/parsimonious dynamics is essential to improve model performance under small datasets. Learning fundamental properties within data, in principle, does not require large or high-coverage datasets. Moreover, a fundamental model can also help remove spurious correlations and maximally promote model stability and generalization under limited data (see the discussion in L54-65 in our paper).\n- T-symmetry happens to be one of the simplest and most fundamental properties that we can leverage to enforce such a fundamental property (see the discussion in L54-65, L104-129). We further extend T-symmetry to make it broadly applicable to generic MDP settings (see the discussion in L113-129).\n- By modeling the latent dynamics to be first-order ODE systems (forcing the latent dynamics to be mathematically simple) and enforcing extended T-symmetry, we can obtain a well-regularized dynamics model (TDM). This model can provide a more effective state-action representation for offline RL (see the discussion in L184-199 and our ablation on representation learning in Fig. 4).\n- The deviation in latent actions and the consistency with T-symmetry specified in TDM actually provide another perspective to detect unreliable or non-generalizable samples. These deviations can serve as a new set of policy constraints to replace the highly restrictive OOD regularizations in existing offline RL algorithms (see the discussion in L200-207, 216-222). Moreover, as we have shown in Fig. 7, this new type of policy constraint allows the policy to reliably generalize to OOD regions without being constrained to the training data distribution, leading to better small-sample performance (see L307-319 for detailed discussion).\n- Lastly, compliance with T-symmetry also enables a reliable latent data augmentation scheme that further addresses the limited size of training data (see empirical analysis in L287-293).\n\n\n> **W2 & Q2. Insufficient discussion on the related works, especially PlayVirtual [1].**\n\nWe thank the reviewer for providing this reference, and we will add it to our final paper. Regarding the relationship and difference between TSRL and PlayVirtual: \n- **Problem setting and model construction:** PlayVirtual is designed for online RL settings where the dynamic models can be continuously improved with fresh interaction data to mitigate multi-step rollout compounding errors. While in offline small dataset settings, multi-step rollouts can lead to considerable compounding errors, even adding cycle-level forward & backward consistency may not be sufficient to properly regulate the model.\nOn the other hand, our proposed TDM **does not perform any rollout generation**. It directly regulates the T-symmetry consistency between ODE latent forward and reverse dynamics at each step ($F(s,a) = -G(s',a)$), which applies a much stronger regulation to improve the model's small-sample performance. \n- **Model learning:** PlayVirtual learns the forward and backward dynamics models with the supervision of the future/previous state representations. While in our TDM, the latent forward and reverse dynamic model are learned as first-order ODEs ($F(s,a)={dz_s \\over dt}$ ; $G(s',a)={-dz_s \\over dt}$ ) to extract the essential patterns within the dynamical system. \n- **Integration with RL algorithm:** PlayVirtual directly incorporate the learned dynamics with standard online RL algorithm (e.g., SAC and Rainbow). While in our paper, we propose a new offline RL algorithm TSRL that closely integrates the learned TDM, allowing the full use of the dynamics-enhanced information from TDM to improve offline RL performance under limited data.\n\n\n> **Q3. Are there any ideas to extend the core idea of this work to online RL algorithms?**\n\nWe thank the reviewer for this thoughtful comment. In principle, our method is also applicable to the online setting, but there are several modifications that could be introduced to our method to enable the best performance:\n- A small set of initial samples need to be collected in the environment using an initial or random policy to warm start the learning of TDM. This will make the learned representation evolve smoothly during the early RL training stage.\n- The strong regularizations in TDM could be properly relaxed, as online samples can be used to continuously improve model learning. We can trade off some regularization to promote model expressiveness.\n- The offline backbone RL algorithm in TSRL needs to be changed to an online RL algorithm with some exploration schemes to remove the pessimism. The T-symmetry regularized policy constraints are probably not necessary for the online setting and can be removed. \n"}, {"review_id": "TW68ZiFSW4", "reviewer": "Reviewer_dZPC", "summary": "The current offline RL algorithm requires a large amount of offline dataset training and has poor performance on small datasets. This article proposes a framework to address this issue. By learning a T-symmetry enhanced dynamic model, capture more fundamental dynamic relationships. Afterward, the article applies T-symmetry to offline RL, uses a model to regularize constraints in the latent action space, and uses a model to filter data for data augmentation. The experimental results showed improvement in small datasets.", "questions": "1. Compared with other methods, since t-symmetry can learn better representation and dynamics, why is there a great improvement in small datasets, but not in full datasets?\n\n2. Why choose to perturb the latent space and then filter the data, rather than perturbing directly on the state space, and then filtering? Is this part of the performance improvement caused by the latent space, the filtered data, or both?\n", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The key idea is novelty. This paper explores symmetries to enhance the performance of offline RL with small datasets. The proposed method has a new technical insight.\n\n2. The article is clearly structured.\n\n3. The proposed method improves performance on offline RL with small datasets.\n", "weaknesses": "1. This article chooses the backward model because of irreversible action, but this does not quite meet the definition of time-reversal symmetry. Other articles (https://arxiv.org/abs/2111.12600) respectively choose the inverse model to get the reversible action, and then process the irreversible action, rather than simply giving up.\n\n2. The proposed method has a great improvement on small datasets, but there is still a lot of gap from the performance of complete data.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **W1. Difference with the origianl definition of time-reversal symmetry. Comparison with the treatment on irreversible actions in https://arxiv.org/abs/2111.12600.**\n\nWe thank the reviewer for providing this reference and will add it to our final paper. Regarding the differences:\n\n- As we have discussed in the last subsection in Section 2 (L113-129), the original T-symmetry (i.e., $F(s,a)=\\dot{s}=-\\tilde{G}(s',a')$) can sometimes be broken by irreversible action $a'$. Hence we consider an extended form, that is, enforcing $F(s,a)=\\dot{s}=-G(s',a)$ while preserving the first-order ODE requirement in T-symmetry. Note that this condition is almost universally held in typical MDP, as it essentially suggests the distribution $P(s,a,s')$ and $P(s',a,s)$ should be equal. It also has the advantage of being irrelevant to the impact of irreversible actions $a'$. Furthermore, using the reverse dynamics model (e.g., function mapping with $(s',a)\\rightarrow s$) has already been adopted in many RL studies [1,2,3], our method differs from these works in that we use it to construct T-symmetry consistency and models it as an ODE system.\n- The CCWM (Cycle-Consistency World Model) mentioned by the reviewer actually adopts a more engineering-oriented approach to address the issue of irreversible actions. It removes \"irreversible\" transitions that produce sudden changes in Q-values over a trajectory in the latent space from the modeling process. This actually ignored some dynamics properties in an MDP associated with such irreversible actions. While in our work, such \"irreversible\" transitions will not cause a problem and can be implicitly captured in our modeling process.\n\n\n> **W2 & Q1. The proposed method has a great improvement on small datasets, but there is still a lot of gap from the performance of complete data.**\n\n- As we have discussed in the conclusion, there is a trade-off between model generalizability and expressiveness. In TDM, we introduce multiple regularizers as well as the additional T-symmetry regularization to obtain a well-behaved dynamics model for the small-dataset setting. But this could hurt model expressiveness, as the model tries to extract the fundamental representations/patterns within the dataset, rather than fits each individual sample. This is beneficial for improving robustness and generalization in the low-data regime but can be overly restrictive for large datasets.\n- In offline RL, learning from very small datasets and large datasets faces different challenges and potentially requires different algorithm design logics. For large datasets with high state-action space coverage, maximally exploiting the offline dataset and ensuring policy learning within data distribution to avoid distributional shift is already sufficient, which is exactly how most existing offline RL algorithms are designed. While under small datasets, strictly regularizing policy within data distribution will hurt performance, encouraging OOD generalization is the key to achieving higher performance. The data scarcity in the small-sample setting makes it more challenging to learn a reliable policy, thus requiring robust regularization techniques to prevent overfitting and promote generalization performance.\n- Finally, it should be noted that although our proposed TSRL is designed for the small-dataset setting, as shown in Table 1 and Fig. 3, it can still achieve comparable (sometimes even better) performance with existing advanced offline RL methods in the complete datasets, despite the fact that D4RL provides overly large amounts of data for simple locomotion tasks.\n\n\n> **Q2. Why choose to perturb the latent space and then filter the data, rather than perturbing directly on the state space, and then filtering?**\n\n- First, in our TDM, the T-symmetry consistency is enforced on the latent forward and reverse dynamics rather than on the original state-action space. Thus it has to evaluate T-symmetry violations in the latent space.\n- The forward and reverse dynamics are learned in latent space because they need to be first-order ODE systems to establish the T-symmetry relationship. For a general nonlinear dynamical system, directly fitting a first-order ODE system can suffer from relatively large errors. Hence the current best practices in the control community (such as Koopman theory [4,5] and Sindy [6,7]) first map the original state-action space into a well-behaved latent space and then construct latent first-order ODE dynamics by fitting the data.\n- Lastly, perturbing in the latent space for data augmentation is much more convenient for implementation. As in our TSRL, the inputs of the Q-function are the latent states and actions ($(z_s, z_a)=\\phi(s,a)$) from the TDM encoder $\\phi$. And the augmented data is executed in the latent space while training the Q function. Perturbing directly on the original state space will incur another origin-to-latent conversion step and cause extra computation.\n\n\n## References\n[1] Lai, H. et al., Bidirectional model-based policy optimization. ICML 2020.\n\n[2] Wang, J. et al., Offline reinforcement learning with reverse model-based imagination. NeurIPS 2021.\n\n[3] Lyu, J., Li, X., & Lu, Z. Double Check Your State Before Trusting It: Confidence-Aware Bidirectional Offline Model-Based Imagination. NeurIPS 2022.\n\n[4] Weissenbacher, M., et al., Koopman qlearning: Offline reinforcement learning via symmetries of dynamics. ICML 2022.\n\n[5] Mezic, I. Spectral properties of dynamical systems, model reduction and decompositions. Nonlinear Dynamics, 2005.\n\n[6] Brunton, S., et al., Discovering governing equations from data by sparse identification of nonlinear dynamical systems. PNAS, 2016.\n\n[7] Champion, K., et al., Data-driven discovery of coordinates and governing equations. PNAS, 2019."}, {"review_id": "JQptz7wH6X", "reviewer": "Reviewer_xNUN", "summary": "This paper proposes 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, embedding and enforcing T-symmetry between a pair of latent forward, and reversing ODE dynamics to learn fundamental dynamics patterns in data. Empirical results on D4RL benchmark datasets validate the good generalization ability of TSRL.", "questions": "- Explain the reason that the performance of TSRL is not comparable with other baselines in most Adroit human and cloned tasks.\n- Add comparisons of other offline reinforcement learning methods using data augmentation.\n- Add related work of the data augmentation methods using offline RL in Sec. 6. ", "limitations": "The limitations of the paper are properly addressed by the authors, but the societal impacts of the paper are not discussed.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The idea is interesting and makes intuitive sense.\n- The paper is overall well-written.\n- Authors do comprehensive experiments on D4RL tasks to evaluate the generalization ability of the new method.", "weaknesses": "- The performance of TSRL is not comparable with other baselines in most Adroit human and cloned tasks.\n- No comparison of other offline reinforcement learning methods using data augmentation.\n- No experiments to validate TSRL alleviates the problem of over-conservatism.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> **W1 & Q1. The performance of TSRL is not comparable with other baselines in most Adroit human and cloned tasks.**\n\n- The full results for Adroit tasks were listed in Appendix C Table 4 due to the space limit of the main article, please check our supplementary material for details. we observe that TSRL achieves much better performance as compared to the baseline algorithms in the pen tasks (both the full datasets and the reduced-size datasets), and comparable performance in other tasks. Note that Adroit tasks are substantially more challenging as compared to MuJoCo tasks due to their high-dimensionality and potentially non-Markovian (human dataset) property, most offline RL algorithms struggle in these tasks.\n- In Appendix C, we also provide the comparative performance of TSRL on Antmaze-umaze tasks with full and 10k reduced-size datasets. We find TSRL also achieves substantially better performance under small datasets of these tasks. \n\n\n\n> **W2 & Q2. No comparison of other offline reinforcement learning methods using data augmentation.**\n\n- In Fig. 5 and discussion in Section 5 L287-293, we have provided the comparison of TSRL with S4RL[1] (reference [42] in the paper) with Gaussian noises on two MuJoCo tasks. \n- To fully address the concern of the reviewer, we have conducted additional experiments on the reduced-size 10k MuJoCo datasets to compare TSRL and other offline RL methods using data augmentation. In particular, we compared with model-free method S4RL[1] with Gaussian (S4RL-N) and uniform (S4RL-U) noises, as well as a recent model-based data augmentation method CABI[2] (Reference [43] in the paper). CABI employs a pair of predictive dynamics models to assess the reliability of the augmented data. The KFC[3] (reference [32] in the paper) is also a model-based data augmentation method, which utilizes Koopman theory to model the dynamical system and augments data points in the latent linearizable space. However, KFC has no publicly available code, and re-implementing this algorithm is quite challenging, given the limited rebuttal period, we opt to evaluate CABI with our method. The detailed results are reported in the following table.\n\n**Table 1: Performance of data augmentation methods with 10k reduced-size D4RL datasets. Each result is generated by 3 random seeds.**\n\n| **D4RL tasks**| **CABI** |**S4RL-N**| **S4RL-U**| **TSRL**|\n| -------|----|----|----|----|\n|Hopper-m|48.3 $\\pm$ 3.9| 28.3 $\\pm$ 6.2| 23.6 $\\pm$ 4.7| **62.0 $\\pm$ 3.7**\n|Hopper-mr|19.8 $\\pm$ 3.9|16.6 $\\pm$ 12.9|12.5 $\\pm$ 12.2|**21.8 $\\pm$ 8.2**\n|Hopper-me|38.3 $\\pm$ 5.3|12.5 $\\pm$ 3.8| 13.1 $\\pm$ 4.7|**50.9 $\\pm$ 8.6**\n|Hopper-e|34.6 $\\pm$ 24.4|14.1 $\\pm$ 12.9|12.2 $\\pm$ 11.6|**82.7 $\\pm$ 21.9**\n|Halfcheetah-m| 34.8 $\\pm$ 1.9|25.1 $\\pm$ 6.8|23.2 $\\pm$ 7.1|**38.4 $\\pm$ 3.1**\n|Halfcheetah-mr| 23.5 $\\pm$ 3.4|15.1 $\\pm$ 9.3|14.8 $\\pm$ 9.5|**28.1 $\\pm$ 3.5**\n|Halfcheetah-me| 29.9 $\\pm$ 1.7|27.1 $\\pm$ 7.1|23.4 $\\pm$ 8.2|**39.9 $\\pm$ 21.1**\n|Halfcheetah-e| 4.2 $\\pm$ 4.1|2.4 $\\pm$ 3.9|1.8 $\\pm$ 3.1|**40.6 $\\pm$ 24.4**\n|Walker2d-m| 42.4 $\\pm$ 23.3|24.5 $\\pm$ 4.3|21.9 $\\pm$ 4.8|**49.7 $\\pm$ 10.6**\n|Walker2d-mr|11.7 $\\pm$ 7.6|1.5 $\\pm$ 2.1|1.4 $\\pm$ 2.3|**26.0 $\\pm$ 11.3**\n|Walker2d-me|17.4 $\\pm$ 9.2|21.9 $\\pm$ 16.4|16.0 $\\pm$ 13.2|**46.7 $\\pm$ 17.4**\n|Walker2d-e|20.2 $\\pm$ 3.4|56.5 $\\pm$ 26.7| 51.1 $\\pm$ 29.7|**102.2 $\\pm$ 11.3**\n \n \n- The results clearly show that TSRL outperforms all offline RL baselines with data augmentation under small datasets. It is also observed that model-based methods TSRL and CABI generally perform better than model-free data augmentation method S4RL in this setting, due to access to additional dynamics information. Moreover, as CABI does not learn a strongly regularized dynamics model with T-symmetry consistency as in our proposed TDM, it still has a noticeable performance gap as compared to our method. \n\n\n> **W3. No experiments to validate TSRL alleviates the problem of over-conservatism.**\n\n- We have conducted experiments to evaluate the OOD generalization performance of TSRL, which is direct evidence to show TSRL can alleviate the problem of over-conservatism. If the reviewer check Fig. 7, L307-319 in Section 5 as well as Fig. 8, L592-607 in Appendix C, we constructed two low-speed datasets from D4RL-Walker2d datasets by removing all high x-velocity samples. This creates smaller datasets with a large proportion of state-action space and transition dynamics unobserved. As the rewards in these tasks encourage high-speed behavior, we want to test if the agent can still generalize and learn well given only these low-speed data. \n- As shown in Fig.7 and 8, the existing offline RL algorithms perform poorly when trained with only low-speed data, which is primarily due to the adoption of over-conservative data-related regularizations, making these algorithms not able to deviate largely from the data distribution. By contrast, 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 is evident if the reviewer inspects the policy rollout distributions (right figures of Fig. 7 and 8) that the policies learned by TSRL indeed generalize to novel high-speed behaviors that are not present in the training data.\n\n\n> **Q3. Add related work of the data augmentation methods using offline RL in Sec. 6.**\n\n- We thank the reviewer's suggestion. We will add more detailed disussion on related work and include the additional experiment results in our final paper.\n\n\n\n## References\n[1] Sinha, S., Mandlekar, A., & Garg, A. S4rl: Surprisingly simple self-supervision for offline reinforcement learning in robotics. CoRL 2022.\n\n[2] Lyu, J., Li, X., & Lu, Z. Double Check Your State Before Trusting It: Confidence-Aware Bidirectional Offline Model-Based Imagination. NeurIPS 2022.\n\n[3] Weissenbacher, M., et al., Koopman q learning: Offline reinforcement learning via symmetries of dynamics. ICML 2022."}]

Author rebuttal:

Meta-review: This paper explores the enhancement of dynamic model learning in offline RL by leveraging the symmetry of forward and reverse dynamics. The submission garnered predominantly favorable feedback. All reviewers concurred on the merit and novelty of the concept. While there were initial reservations regarding the robustness of the experiments, these concerns were substantially mitigated following the rebuttal and subsequent discussions. Consequently, the AC has decided to accept the submission.
Decision: Accept (poster)