Title: Look Beneath the Surface: Exploiting Fundamental Symmetry for Sample-Efficient Offline RL
Venue: NeurIPS 2023
Accepted: yes
Keywords: sample efficiency; offline reinforcement learning; fundamental symmetry
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

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

1 Introduction

The recently emerged offline reinforcement learning (RL) provides a new paradigm to learn policies from pre-collected datasets without the need to interact with the environments [1, 2, 3]. This is particularly desirable for solving practical tasks, as interacting with real-world systems can be costly or risky, and high-fidelity simulators are also hard to build [4]. However, existing offline RL methods have high requirements on the size and quality of offline datasets in order to achieve reasonable performance. When such requirements are not met, these algorithms may suffer from severe performance drop, as illustrated in Figure 1. Current offline RL algorithms are trained and validated on benchmark datasets (e.g., D4RL [5]) that contain millions of transitions for simple tasks. Whereas under realistic settings, it is often impractical or costly to collect such a large amount of data, and the real datasets might only narrowly cover the state-action space. Clearly, learning reliable policies from

small datasets with partial coverage has become one of the most pressing challenges for successful real-world deployments of offline RL.

Unfortunately, sample-efficient design considerations have been largely overlooked in the majority of offline RL studies. Pessimism is universally adopted in existing offline RL methods and various forms of data-related regularizations have been applied to combat the distributional shift and exploitation error accumulation issues [3, 6], such as conservatively restricting policy deviation from the behavioral data [3, 6, 7, 8], regularizing value function on outof-distribution (OOD) samples [9, 10, 11, 12], learning policy on a pessimistic MDP [4, 13, 14],

walker2d-medium-v2

walker2d-expert-v2

120

120

10k

10k

100

100

Normalized scores

Normalizedscores

80

80

60

60

40

40

20

20

0

BCTD3+BC MOPO

CQL

IQL

TSRL

BC

TD3+BC MOPO

CQL

IQL

TSRL

or adopting strict in-sample learning [15, 16, 17, 18]. In many of these methods, the full coverage assumption plays an important role in their theoretical performance guarantees [3, 19], which assumes the dataset to contain all state-action pairs in the induced distribution of the policy. Obviously, most of the state-action space will become OOD regions under a small dataset. Applying strict data-related regularizations will inevitably cause severe performance degradation and poor generalization [20]. Consequently, it is important to rethink what is essential in policy learning with small datasets. In other words, what is the fundamental or invariant information that can be used to facilitate policy learning, without being conservatively confined by the limited data?

In this paper, we provide a new insight that exploiting the fundamental symmetries in the system dynamics can substantially enhance the performance of offline RL with small datasets. Specifically, we consider the time-reversal symmetry (also called T-symmetry ), which is one of the most fundamental properties discovered in classical and quantum mechanics [21, 22]. It suggests that the underlying laws of physics should not change under the time-reversal transformation: t → -t [22, 23]. Specifically, we are interested in an extended form of T-symmetry for MDP due to its simplicity and universality in physical systems. In the small dataset setting, enforcing T-symmetry in dynamics model learning and offline RL offers three crucial benefits. First, as T-symmetry describes a fundamental property of a system, hence learning it in principle would not require a large or high-coverage dataset. Second, T-symmetry captures what is essential and invariant in the dynamical system, thus offering great generalizability. Lastly, compliance with T-symmetry also provides an important clue to detecting unreliable or non-generalizable state-action samples.

Based on these intuitions, we develop a physics-informed T-symmetry enforced Dynamics Model ( TDM ) to learn a well-behaved and generalizable dynamics model with small datasets. TDM enforces the extended T-symmetry between a pair of latent space forward and reverse dynamics sub-models, which are modeled as first-order ordinary differential equation (ODE) systems to extract fundamental dynamics patterns in data. TDM provides both well-behaved representations for small datasets and a new reliability measure for OOD samples based on compliance with the T-symmetry. These can be used to construct a highly data-efficient offline RL algorithm, which we call T-Symmetry regularized offline RL ( TSRL ). Specifically, TSRL uses the T-symmetry regularized representations learned in TDM to facilitate value function learning. Furthermore, the deviation on latent actions and the consistency with T-symmetry specified in TDM provide another perspective to detect unreliable or non-generalizable samples, which can serve as a new set of policy constraints to replace the highly restrictive OOD regularizations in existing offline RL algorithms. Lastly, a reliable latent space data augmentation scheme based on compliance with the T-symmetry is also applied to further remedy the limited size of training data. With these designs, TSRL performs surprisingly well compared with the state-of-the-art offline RL algorithms on reduced-size D4RL benchmark datasets with even as few as 1% of the original samples. To the best of the authors' knowledge, this is the first offline RL method that demonstrates promising performance on extremely small datasets.

2 Preliminaries

Offline reinforcement learning. We consider the standard Markov decision process (MDP) setting [24], which is represented as a tuple M = {S , A , r, P , ρ, γ } , where S and A are the state and

action spaces, r ( s, a ) is a scalar reward function, P is the transition dynamics, ρ is the initial state distribution, and γ ∈ (0 , 1) is a discount factor. The objective of RL is to learn a policy π ( a | s ) by maximizing the expected cumulative discounted return E π [ ∑ ∞ t =0 γ t r ( s t , a t ))] , which is typically approximated by a value function Q ( s, a ) using some function approximators, such as deep neural networks. The Q-function is typically learned by minimizing the squared Bellman error:

where ˆ Q denotes a target Q-function, which is a delayed copy of the current Q-function; B π is the Bellman operator, which is often used as the Bellman evaluation operator B π ˆ Q ( s, a ) = r ( s, a ) + γ E a ′ ∼ π ˆ Q ( s ′ , a ′ ) in many RL algorithms.

Under the offline RL setting, we are provided with a fixed dataset D = { ( s 0 , a 0 , r 0 , s 1 , · · · ) ( i ) } N i =1 without any chance of further environment interactions. Directly applying standard online RL methods in the offline setting suffers from severe value overestimation, due to counterfactual queries on OOD data and the resulting extrapolation errors [1, 3, 6]. To avoid this issue, a widely used offline RL framework adopts the following behavior regularization scheme which regularizes the divergence between the learned policy π and the behavior policy π β of the dataset D :

where D ( ·∥· ) is some divergence measures, which can have either an explicit [3, 8] or implicit form [6, 17, 25]. Although straightforward, existing behavior regularization methods have been shown to be over-conservative [9, 20] due to the restrictive regularization with respect to the behavior policy in data, which may suffer from notable performance drop under small datasets.

Time-reversal symmetry in dynamical systems. Most real-world dynamical systems with state measurement x ∈ Ω on some phase space Ω can be modeled or approximated by the system of non-linear first-order ordinary differential equations (ODEs) as d x dt = F ( x ) , where F is some general non-linear, at least C 1 -differentiable vector-valued function. First-order ODE systems are said to be time-reversal symmetric if there is an invertible transformation Γ : Ω ↦→ Ω , that reverses the direction of time [22, 26]: d Γ( x ) /dt = -F (Γ( x )) . If we define a time evolution operator U ∆ t : Ω ↦→ Ω as U ∆ t : x ( t ) ↦→ U ∆ t ( x ( t )) = x ( t +∆ t ) . Then T-symmetry implies that Γ ◦ U τ = U -τ ◦ Γ . 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 = ds dt as the time-derivative of the current state s , which can be approximated as the difference between the next and current states, i.e., ˙ s = s ′ -s . For a dynamical system that satisfies T-symmetry, it suggests that if we learn a forward dynamics F ( s, a ) = ˙ s and a reverse dynamics ˜ G ( s ′ , a ′ ) = -˙ s as a pair of first-order ODEs, we should have F ( s, a ) = -˜ G ( s ′ , a ′ ) .

However, from a decision-making perspective, it is known that T-symmetry can sometimes be broken by irreversible actions or some special dynamic processes (e.g., frictional force against motion). Hence in this paper, we consider a more generic treatment by leveraging an alternative ODE reverse dynamics model G ( s ′ , a ) = -˙ s to establish the T-symmetry with the forward dynamics, i.e., enforcing F ( s, a ) = -G ( s ′ , a ) . Note that G ( s ′ , a ) is now defined on the next state s ′ and the current action a , rather than the next action a ′ , thus is not impacted if the next action is irreversible. This extended T-symmetry provides a more fundamental and almost universally held property in discrete-time MDP systems. Its simplicity and fundamentalness make it an ideal property that we can leverage to construct a well-behaved data-driven dynamics model and a robust offline RL algorithm under small datasets.

3 T-Symmetry Enforced Dynamics Model

In this section, we present the detailed design of TDM, which is capable of learning a more fundamental and T-symmetry preserving dynamics from small datasets. The key ingredients of TDM are to embed a pair of latent forward and reverse dynamics as ODE systems, and further enforce their T-symmetry consistency. This design offers several benefits. First, embedding ODE modeling

VsZs·S

Encoder

Decoders

Reconstruction:rec=Ils-s(zs,0)l2+Ila-a(za)2

z →s decoding: &as = I/s - s(f(Φ(s,a), 1)2

s(Zs,0)→

S

Latentforwarddynamics

f

s

s(zs,1)→

Forward ODE:

(sa) s - f((s, a)2

as

(sa)(s) - g(中(s',a)

Latentreversedynamics

Reverse ODE:

&fwd=

ds'

g

4a

→ya(za)→á

T-symmetry:

-sym =Ig(zs+f(zsza),za)+f(zs,za)l2

T-symmetry

LTDM=

[&rec + &as + &fwa + &rus + r-sym] + A[<(f) + L(g)]

VsZs'·(-s)

(s,a,S)ED

of the system helps to extract more fundamental descriptions of how the system dynamics evolve over time in the latent space. Similar ideas have also been explored in Koopman theory [27, 28] and discovering governing equations of nonlinear dynamical systems [29, 30]. The requirement on extended T-symmetry for both forward and reverse dynamics introduces additional consistency regularization, which is helpful to improve model learning stability under limited data. With a regularized and well-behaved dynamics characterization, we can learn more effective state-action representations, which allows removing non-essential elements in the raw data to promote generalization performance. Lastly, compliance with T-symmetry can provide extra information regarding the reliability of OOD samples, which can be particularly useful in downstream offline policy learning. As illustrated in Figure 2, the proposed TDM consists the following components:

Encoder and decoders. TDM implements a state-action encoder ϕ ( s, a ) = ( z s , z a ) and a pair of decoders ψ s ( · , δ s ) , ψ a ( z a ) = a that embed the state-action pair ( s, a ) into latent representations ( z s , z a ) and then map them back. Specifically, we require the state decoder ψ s ( · , δ s ) to be capable of decoding both z s and ˙ z s , where δ s is an indicator to help the decoder to decide the target output, with δ s = 0 as decoding z s → s and δ s = 1 as decoding ˙ z s → ˙ s . The encoder ϕ and decoders ψ s and ψ a induce the following reconstruction loss term for each state-action pair ( s, a ) :

Latent forward dynamics. Inspired by the prior works that incorporate physics-informed information into dynamical systems modeling [27, 29, 30], we embed a discrete-time first-order ODE system to capture the latent forward dynamics f ( z s , z a ) = ˙ z s . Similar to ˙ s , we write ˙ z s = z s ′ -z s to denote the forward difference of the next and current latent state representations. Note that based on the chain-rule, we have ˙ z s = dz s dt = ∂z s ∂s · ds dt = ∇ s z s · ˙ s . To enforce the ODE property, we can introduce the following loss term for f :

Minimizing L fwd ensures that the latent forward dynamics f correctly predicts the forward time evolution of latent states in the dynamical system. We also require the decoder ψ s ( · , δ ˙ s ) to be able to decode ˙ s from ˙ z s to ensure it is compatible with the ODE property, which implies the following loss:

Latent reverse dynamics. We can further introduce a latent reverse dynamics g ( z s ′ , z a ) = -˙ z s in the model, which captures the reverse time evolution of the system in the latent space. Similar to the forward dynamics loss L fwd , we can write the reverse dynamics loss for g as:

T-symmetry regularization. The above latent forward and reverse dynamics f and g are learned to be two models, which may not necessarily satisfy the proposed extended T-symmetry. We can enforce the extended T-symmetry by requiring f ( z s , z a ) = -g ( z s ′ , z a ) . To further couple the learning process of f and g , note that z s ′ = z s + ˙ z s = z s + f ( z s , z a ) , which suggests g ( z s ′ , z a ) = g ( z s + f ( z s , z a ) , z a ) = -˙ z s = -f ( z s , z a ) . This implies the following T-symmetry consistency loss:

Above instance-wise T-symmetry consistency loss also provides an alternative measure for evaluating the reliability of a data sample. A state-action pair ( s, a ) with a large ℓ T -sym ( ϕ ( s, a )) implies that this sample may not be well-explained by TDM or consistent with the fundamental symmetry of the system. This can be used to detect unreliable OOD samples in offline policy optimization as well as construct a new latent space data augmentation procedure, which will be discussed in later content.

Final learning objective of TDM. Finally, we can formulate the overall loss function of TDM as:

where L 1 ( f ) and L 1 ( g ) are L1-norms of the parameters of f and g , and λ L 1 is a scale parameter. L1 regularization is introduced to encourage learning parsimonious latent dynamics for f and g , which helps to improve model generalizability [29, 30].

Note that the proposed TDM is very different from the conventional dynamics models used in modelbased RL (MBRL) methods [4, 13, 14, 31, 32, 33, 34, 35]. The dynamics models in MBRL focus on constructing a predictive model to represent the forward transition dynamics of the system. Whereas, TDM is formulated as a reconstruction model with T-symmetry preserving embedded ODE latent dynamics, which aims at explaining and extracting the fundamental dynamics of the system. As a result, TDM can be substantially more well-behaved and robust when learning from small datasets.

4 T-Symmetry Regularized Offline RL

In this section, we discuss how to incorporate the properties of TDM to construct a sample-efficient offline RL algorithm, which we call T-Symmetry regularized offline RL (TSRL). Specifically, we leverage three types of information from TDM to jointly improve offline RL performance under small datasets. First, we use the well-behaved state-action representations provided by TDM to facilitate value function learning. Second, the consistency with the latent dynamics in TDM also provides new forms of policy constraints to penalize unreliable and non-generalizable OOD samples. Finally, compliance with T-symmetry enables a new latent space data augmentation procedure to further enhance the algorithm's performance with limited data.

T-symmetry regularized representation. Representation learning has been shown to be an effective approach to enhancing sample efficiency and generalization in many online and offline RL studies [36, 37, 38, 39, 40]. A notable property of TDM is that the learned latent state-action representations from the encoder ( z s , z a ) = ϕ ( s, a ) are compatible with both the latent forward and reverse ODE dynamics f and g . This leads to well-regularized and T-symmetry preserving representations that can potentially generalize better on OOD areas under small dataset settings. We can simply use the latent state-action representation ( z s , z a ) extracted by the encoder ϕ ( s, a ) of TDM in the value function learning, which gives the following policy evaluation objective:

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 [41] and the system dynamics properties. Moreover, restricting policy optimization only within data-covered regions can be overconservative when the offline dataset is small, which can result in degraded policy performance [20]. In TSRL, we instead consider an alternative regularization scheme, which restricts the deviation on latent actions and the T-symmetry consistency of policy-induced samples, corresponding to the following policy optimization objective:

where latent actions z a and z a π are obtained from ϕ ( s, a ) and ϕ ( s, π ( ·| s )) respectively. The second term restricts the latent action z a π of policy π from deviating too much from the latent action z a in data. The third term regularizes the T-symmetry consistency of policy-induced samples ( s, π ( ·| s )) , which is evaluated based on Eq. (7) and the learned TDM. λ 1 and λ 2 are weight parameters, which only need to be roughly adjusted to ensure both the regularization terms are in a similar scale

as the first term. We also introduce a normalization term α on the value function for training stability similar to TD3+BC [8], which is computed based on a training batch B of samples as α 0 / [ ∑ ( s,a ) ∈ B Q ( ϕ ( s, π ( ·| s )))] . We set α 0 = 2 . 5 in all of our experiments without tuning.

Instead of strictly regularizing OOD actions as in existing offline RL algorithms, our policy regularization scheme actually allows some reliable OOD action for policy optimization. As TDM is designed to capture the fundamental and invariant system dynamics patterns in the latent space, if an OOD action has a similar latent representation to some latent actions in data and also agrees with the T-symmetry property in TDM, then we can expect some degree of equivalency between these actions. This leads to more relaxed policy constraints by enabling policy learning and generalization on reliable OOD regions, which is critical for the small dataset setting.

T-symmetry consistent latent space data augmentation. It has been shown in previous studies [42, 28, 43] that data augmentation can potentially improve the function approximation of the Q-networks by smoothing out the learned state-action space, hence often lead to more robust policy and better data efficiency. However, existing data augmentation methods in offline RL studies either blindly add random perturbations to states [42] or utilize costly non-linear symmetry transformations, such as Koopman theory [28]. With TDM, we can provide a very simple yet principled data augmentation scheme based on the T-symmetry property. Assuming we add a small perturbation ϵ to a latent state z s , i.e., ( z s , z a ) ↦→ ( z s + ϵ, z a ) , then the corresponding perturbation ϵ ′ on the next latent state z s ′ according to the latent forward dynamics ˙ z s = f ( z s , z a ) satisfies: z s ′ + ϵ ′ = z s + ϵ + f ( z s + ϵ, z a ) . On the other hand, by the T-symmetry construction in TDM, we can recover back the current perturbed latent state based on the latent reverse dynamics -˙ z s = g ( z s ′ , z a ) as: z s + ϵ ′′ = z s ′ + ϵ ′ + g ( z s ′ + ϵ ′ , z a ) . Clearly, we should have ϵ = ϵ ′′ , which suggests the following condition:

This is exactly equivalent to requiring the instance-wise T-symmetry consistency loss (Eq. (7)) ℓ T -sym ( z s + ϵ, z a ) = 0 . Hence we can use T-symmetry consistency loss ℓ T -sym ( · ) as a reliability measure to filter out unreliable augmented samples ( z s + ϵ, z a ) that are inconsistent with the Tsymmetry property of the learned latent dynamics in TDM. In our implementation, we only keep augmented samples that satisfy ℓ T -sym ( z s + ϵ, z a ) ≤ h , where we consider a non-parametric treatment for threshold h , by setting it as the τ -quantile value of all ℓ T -sym ( ϕ ( s, 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 (10), as well as the T-symmetry consistent latent space data augmentation. In our experiments reported in the next section, we generate K = 1 augmented samples for each transition in the dataset, and filter based on the T-symmetry consistency loss. The pseudo-code of TSRL is summarized in Algorithm 1.

Algorithm 1 T-Symmetry Regularized Offline RL (TSRL)

Require: Offline dataset D , encoder ϕ , latent forward and reverse dynamics models f and g from TDM trained using objective Eq. (8).

1: Compute the T-symmetry consistency loss ℓ T -sym ( ϕ ( s, a )) (Eq. (7)) for all samples in D , and set their τ -quantile value as the augmentation threshold h .

2: Initialize the policy network π , critic networks Q and their target network.

3: for t = 1 , · · · , M training steps do

4: Sample a mini-batch B of samples { ( s, a, r, s ′ ) } ∼ D and compute their representations { ( z s , z a , z s ′ ) } .

5: / / T-symmetry consistent latent space data augmentation

6: Generate K perturbed samples by adding perturbations ϵ ∼ N (0 , 0 . 01 σ z s ) on latent states z s of each sample in B , where σ z s is the std of latent states in data.

7: Add augmented samples ( z s + ϵ, z a , z s ′ + ϵ ′ ) to B if satisfies ℓ T -sym ( z s + ϵ, z a ) ≤ h .

8: / / Critic training with T-symmetry regularized representation

9: Update the value function Q based on the policy evaluation objective Eq. (9).

10: / / Policy training with T-symmetry regularized policy constraints

11: Update the policy π based on the policy improvement objective Eq. (10).

12: Soft update the target networks.

13: end for

5 Experiments

We evaluate TSRL on the D4RL MuJoCo-v2 and Adroit-v1 benchmark datasets [5] against behavior cloning (BC) as well as state-of-the-art (SOTA) offline RL methods, including model-free methods TD3+BC [8], CQL [9], IQL [15], DOGE [20], and model-based method MOPO [13]. In particular, we compare with DOGE [20], which leverages a state-conditioned distance function as policy constraint and exhibits strong OOD generalization performance. We report the final normalized performance of each algorithm after training 1M steps.

Performance on small datasets. We compare the performance of TSRL and the baseline methods on both the full D4RL datasets and their reduced-size datasets with only 5k ∼ 10k samples, which are constructed by randomly sampling a given fraction of trajectories in the full datasets * . These reducedsize 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 [20] is the only baseline that can still achieve reasonable performance in some tasks, primarily due to its capability to leverage the relationship between dataset geometry and the generalization capability of deep Q-functions, which is also beneficial for policy learning under small datasets. However, we still observe that DOGE struggles in some datasets, especially those with an extremely narrow distribution.

By contrast, TSRL achieves substantially better performance in all small dataset tasks, indicating superior sample efficiency. Moreover, although MOPO [13] also learns a dynamics model for offline policy learning, it performs badly when the dataset is small, revealing the importance of using a well-regularized model like TDM in the small-sample regime. It is interesting to see that most

scores

walker2d-medium-v2

Normalizedscores

walker2d-medium-expert-v2

Normalized scores

walker2d-expert-v2

100

80

100

100

Normalized

80

80

60

60

60 =

40

40

40

20

20

BCTD3+BCMOPO

CQL

IQL

DOGE

TSRL

BC

TD3+BCMOPO

CQL

IQL

DOGE

TSRL

0

BCTD3+BCMOPO

CQL

IQL

DOGE

TSRL

1M

100k

50k

10k

5k

Normalizedscores

60

40

20

0

walker2d-medium-v2

walker2d-medium-expert-v2

AE-rep

SimSiam

AE-fwd-rep

TDM-no-ODE

TSRL-nO-A

methods fail on the 10k Halfcheetah-expert dataset, with only TSRL showing some success. We suspect this is due to the extremely narrow data distribution of the reduced dataset, which requires strong OOD generalization ability for algorithms to solve the task.

To further examine the impact of the training data size on algorithm performance, we also conduct experiments on three Walker2d datasets (medium, medium-expert, and expert) by varying the size of samples from 1M to 5k. The results are presented in Figure 3. It can be observed that most baseline offline RL algorithms experience a sharp performance drop when the datasets are reduced to 10k samples. Whereas, TSRL is still capable of preserving reasonable performance as the decrease of data size, even for extremely small datasets that contain only 5k samples.

Investigation on learned representations. To investigate the quality of the latent representation learned in TDM, we compare the performance of different representation learning approaches on the 10k datasets in Figure 4. To solely evaluate the impact of the representation, we remove the latent space data augmentation component from TSRL ('TSRL-no-A') and replace the state-action encoder ϕ ( s, a ) learned from other representation learning approaches, including the autoencoder ('AE-rep'), autoencoder with latent forward dynamics ('AE-fwd-rep') without the ODE structure and the T-symmetry regularization in TDM, and a recent popular self-supervised representation learning method SimSiam [44] ('SimSiam'). To further investigate the impact of enforcing the ODE property, we also consider a variant of TDM ('TDM-no-ODE') by removing the ODE structure in latent forward and reverse dynamics. More detailed experiment setups are presented in Appendix B.

The results demonstrate that TDM representation achieves the best performance in all small-dataset experiments. By comparing 'TSRL-no-A' and 'TDM-no-ODE' with 'AE-rep' and 'AE-fwd-rep', we can see that the bi-directional design and the T-symmetry regularization are crucial for performance improvement. Moreover, we find 'TSRL-no-A' consistently achieves better performance and lower variance as compared to 'TDM-no-ODE', further confirming the benefit of incorporating ODE structure in producing a well-behaved representation for downstream tasks under small datasets.

Evaluation on data augmentation. We evaluate the impact of the proposed T-symmetry consistent latent space data augmentation in Figure 5. The results show that our proposed data augmentation scheme can help speed up convergence and reduce variance. Compared with less principled data augmentation methods such as adding zero-mean Gaussian noises as in S4RL [42], our method offers much better performance improvement. As shown in Figure 5, blindly adding random perturbations could suffer from performance degradation over the course of training, while TSRL with T-symmetry consistent data augmentation enjoys better training robustness. Additional comparative experiments on the proposed data augmentation method can be found in Table 6 of Appendix C.

Additional ablations. We also investigate the joint impact of the three design components in TSRL, including T-symmetry regularized representation and policy constraints, as well as the T-symmetry

NormalizedReturn

40

Normalized Return

40

30

30

20-

20

W

10

10

8.00

8

Training Steps (x106)

0.25

0.50

0.75

1.00

Training Steps (x106)

0.25

0.50

0.75

1.00

TSRL+S4RL

TSRL-no-A

TSRL

Normalized Return

40

Normalized Return

40

30

MM

W

30

20

20

10

10

8.00

0

Training Steps (x106)

0.25

0.50

0.75

1.00

8.00

Training Steps (x106)

0.25

0.50

0.75

1.00

TD3+BC

TSRL-no-A&R

TSRL-no-A&P

TSRL-no-A

TSRL

Walker2d-medium-v2

Distributionofx-velocityinpolicyrollouts

80

Normalized Return

1.4

Training data

1.2-

1.0

Hrv

0.2

0.0

0.25

0.50

0.75

1.00

5.0

2.5

0.0

2.5

5.0

7.5

Walker2d-medium-expert-v2

Distributionofx-velocityinpolicyrollouts

Normalized Return

80

Probability Density

1.2

Training data

.0

0.2Xmax-x-velocity

0.8

0.6-

0.4

0.2

0.05.0

0.25

0.50

0.75

1.00

2.5

0.0

2.5

5.0

7.5

Training Steps (x106)

X-velocity of thetop

TD3+BC

CQL

IQL

TSRL

consistent latent space data augmentation. We include TD3+BC for comparison, as it can be perceived as a vanilla version of TSRL without the previous three components. Moreover, the variant of TSRL with only the latent representation removed is not evaluated, as the latent space data augmentation also depends on the representation provided by TDM. Figure 6 presents the performance of all variants of TSRL on the 10k Halfcheetah medium and medium-expert datasets. As expected, it is observed that T-symmetry regularized representation and policy constraints jointly play a critical role in maintaining the performance of TSRL. As discussed previously, adding T-symmetry consistent latent space data augmentation also shows a positive impact on performance. Furthermore, it is observed that TD3+BC suffers from over-fitting on small datasets as its performance drops significantly with the increase of training steps, especially in the 10k Halfcheetah-medium dataset, while this phenomenon is not observed in TSRL. Nevertheless, the complete TSRL achieves the best performance in both tasks.

Generalization performance. To further verify the generalizability of TSRL, we construct lowspeed datasets from the Walker2d-medium and Walker2d-medium-expert datasets by filtering out all high x-velocity samples (x-velocity of the top > 0 . 2 × max-x-velocity). This results in two smaller datasets with a large proportion of transition dynamics unobserved. We want to test if the agent can still generalize and learn well given only these low x-velocity data, with all high-speed samples removed. The experiment results are presented in Figure 7.

Under this setting, we observed that all existing offline RL baselines perform poorly when trained with only the low-speed dataset. This is primarily due to over-conservative data-related regularizations, which cause ineffective policy learning if the OOD region occupies the majority of state-action

space. CQL performs especially poorly in both tasks, perhaps due to over-conservative value function learning that impedes the policy to acquire some necessary control strategy to finish the task. TD3+BC performs poorly in the low-speed medium dataset, likely because this dataset has a narrower data distribution than the medium-expert dataset, and the latter could still contain some samples with reasonable speeds after filtering with x-velocity > 0 . 2 × max-x-velocity in the dataset. IQL exhibits some level of generalization capability, but is still much weaker as compared to TSRL. By contrast, we observe that TSRL is still able to achieve good performance, due to the access to more fundamental dynamics information that remains invariant in both low- and high-speed data. This can be further verified if we inspect the policy rollout distributions (right figures of Figure 7) that the policies learned by TSRL indeed generalizes to high-speed behavior that is not present in the training data.

6 Related Work

Learning fundamental dynamics in physical systems. Learning conservation laws or invariant properties within a physical system is an active research area in physics [45, 46, 29, 30], climate science [47], and neuroscience [48], etc. A classic approach is based on Koopman theory, which represents the nonlinear dynamics in terms of an infinitedimensional linear operator [27]. In practice, this is achieved by finding a coordinate transformation to produce a finite-dimensional representation in which the non-linear dynamics are approximately linear. However, it also suffers from computationally expensive coordinate transformations and is only able to approximate the system dynamics. Another approach is utilizing a sparse regression model with the fewest terms to describe the nonlinear system dynamics [29, 30]. However, it assumes that the dynamical systems only have a few critical terms, which severely limits the model expressiveness and often requires prior knowledge of these critical terms. Based on expressive deep neural networks, a recently emerged research direction is to build ODE networks to learn conservation law in the dynamical system from data [49, 50, 51, 26]. Our proposed TDM falls within this direction, which models both forward and reverse latent ODE dynamics with deep neural networks and incorporates additional regularization on T-symmetry.

Offline reinforcement learning. Offline RL addresses the challenge of deriving policies from fixed, pre-collected datasets without interaction with the environment. Under this offline learning paradigm, conventional off-policy RL approaches are prone to substantial value overestimation when there is a large deviation between the policy and data distributions. Existing offline RL methods address this issue by following several directions, such as constraining the learned policy to be 'close' to the behavior policy [6, 3, 8, 52], regularizing value function on OOD samples [9, 10, 11, 53], enforcing strict in-sample learning [16, 15, 17, 18, 54], and performing pessimistic policy learning with uncertainty-based reward or value penalties [13, 14, 4, 12, 55, 33]. Most existing offline RL methods adopt the pessimism principle and avoid policy evaluation on OOD samples. Although this treatment helps to alleviate exploitation error accumulation, it can be over-conservative and causes severe performance degradation if the training dataset is small or has poor state-action space coverage [20]. TSRL tackles this issue by allowing dynamics explainable OOD samples for policy optimization, thus offering greatly improved small-sample performance.

7 Discussion and Conclusion

In this paper, we propose a physics-informed dynamics model TDM and a new offline RL algorithm TSRL, which exploit the fundamental symmetries in the system dynamics for sample-efficient offline policy learning. TDM embeds and enforces T-symmetry between a pair of latent forward and reverse ODE dynamics to learn fundamental dynamics patterns in data. The well-behaved representations and a new reliability measure for OOD samples based on T-symmetry from TDM can be readily used to construct the proposed TSRL algorithm, which achieves strong performance on small D4RL benchmark datasets and exhibits good generalization ability. There are also some limitations in our proposed approach. For example, in order to learn a well-behaved dynamics model, we introduced a set of dynamics and symmetry regularizations in TDM, which are beneficial to improve model generalization, but will lose some model expressiveness. However, we believe this can be a worthwhile trade-off between precision and generalization under small dataset settings, due to substantially improved model robustness.

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

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.

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.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Systems , pages 11761-11771, 2019.

Xianyuan Zhan, Haoran Xu, Yue Zhang, Xiangyu Zhu, Honglei Yin, and Yu Zheng. Deepthermal: Combustion optimization for thermal power generating units using offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 4680-4688, 2022.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning , pages 2052-2062. PMLR, 2019.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in Neural Information Processing Systems , 34, 2021.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Neural Information Processing Systems (NeurIPS) , 2020.

Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference on Machine Learning , pages 5774-5783. PMLR, 2021.

Haoran Xu, Xianyuan Zhan, and Xiangyu Zhu. Constraints penalized q-learning for safe offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 8753-8760, 2022.

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.

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. In Neural Information Processing Systems (NeurIPS) , 2020.

Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. In Neural Information Processing Systems (NeurIPS) , 2020.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In International Conference on Learning Representations , 2022.

David Brandfonbrener, Will Whitney, Rajesh Ranganath, and Joan Bruna. Offline rl without off-policy evaluation. Advances in Neural Information Processing Systems , 34:4933-4946, 2021.

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.

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:40854098, 2022.

Jinglin Chen and Nan Jiang. Information-theoretic considerations in batch reinforcement learning. In International Conference on Machine Learning , pages 1042-1051. PMLR, 2019.

Jianxiong Li, Xianyuan Zhan, Haoran Xu, Xiangyu Zhu, Jingjing Liu, and Ya-Qin Zhang. When data geometry meets deep function: Generalizing offline reinforcement learning. In The Eleventh International Conference on Learning Representations , 2022.

JP Elliott and PG Dawber. Symmetry in Physics, Volume 1 and 2 . Macmillan, 1979.

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.

George W Bluman and Sukeyuki Kumei. Symmetries and differential equations , volume 81. Springer Science & Business Media, 2013.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

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.

In Huh, Eunho Yang, Sung Ju Hwang, and Jinwoo Shin. Time-reversal symmetric ode network. Advances in Neural Information Processing Systems , 33:19016-19027, 2020.

Igor Mezi´ c. Spectral properties of dynamical systems, model reduction and decompositions. Nonlinear Dynamics , 41:309-325, 2005.

Matthias Weissenbacher, Samarth Sinha, Animesh Garg, and Kawahara Yoshinobu. Koopman q-learning: Offline reinforcement learning via symmetries of dynamics. In International Conference on Machine Learning , pages 23645-23667. PMLR, 2022.

Steven L Brunton, Joshua L Proctor, and J Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the national academy of sciences , 113(15):3932-3937, 2016.

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 , 116(45):22445-22451, 2019.

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.

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.

Xianyuan Zhan, Xiangyu Zhu, and Haoran Xu. Model-based offline planning with trajectory pruning. In Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence, IJCAI-22 , pages 3716-3722, 7 2022.

Tao Yu, Cuiling Lan, Wenjun Zeng, Mingxiao Feng, Zhizheng Zhang, and Zhibo Chen. Playvirtual: Augmenting cycle-consistent virtual trajectories for reinforcement learning. Advances in Neural Information Processing Systems , 34, 2021.

Changmin Yu, Dong Li, Jianye Hao, Jun Wang, and Neil Burgess. Learning state representations via retracing in reinforcement learning. arXiv preprint arXiv:2111.12600 , 2021.

Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. 2021.

Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In International Conference on Machine Learning , pages 5639-5650. PMLR, 2020.

Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Marc G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference on Learning Representations , 2021.

Mengjiao Yang and Ofir Nachum. Representation matters: offline pretraining for sequential decision making. In International Conference on Machine Learning , pages 11784-11794. PMLR, 2021.

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.

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.

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.

Jiafei Lyu, Xiu Li, and Zongqing Lu. Double check your state before trusting it: Confidenceaware bidirectional offline model-based imagination. Advances in Neural Information Processing Systems , 35:38218-38231, 2022.

Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 15750-15758, 2021.

John David Anderson and J Wendt. Computational fluid dynamics , volume 206. Springer, 1995.

Ilia Grigorenko and Elena Grigorenko. Chaotic dynamics of the fractional lorenz system. Physical review letters , 91(3):034101, 2003.

Kevin E Trenberth and Kevin E Trenberth. Climate system modeling . Cambridge University Press, 1992.

Eugene M Izhikevich. Dynamical systems in neuroscience . MIT press, 2007.

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018.

Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural odes. Advances in Neural Information Processing Systems , 32, 2019.

Xuanqing 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.

Haoran Xu, Xianyuan Zhan, Jianxiong Li, and Honglei Yin. Offline reinforcement learning with soft behavior regularization. arXiv preprint arXiv:2110.07395 , 2021.

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-36612, 2022.

Xiangsen Wang, Haoran Xu, Yinan Zheng, and Xianyuan Zhan. Offline multi-agent reinforcement learning with implicit global-to-local value regularization. Advances in Neural Information Processing Systems , 2023.

Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. Advances in neural information processing systems , 34:7436-7447, 2021.

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research , 15(1):1929-1958, 2014.

A Implementation Details

Implementation details for TDM. As discussed in Section 3 and the illustrative Figure 2, TDM is modeled as a physics-informed reconstruction model with embedded ODE latent dynamics and T-symmetry preserving design.To ensure optimal training performance of TDM, we have included some additional implementation details below. More hyperparameters details of TDM are discussed in the next section.

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 ψ ( · , δ s ) , we concatenate an extra indicator δ s in the input to help the state decoder to decide the target output. More specifically, to decode z s → s , we concatenate δ s = 0 with z s as input; and for ˙ z s → ˙ s , we concatenate δ s = 1 with ˙ z s .

Computing second derivative of ϕ ( · ) : As TDM involves a pair of latent ODE forward and reverse dynamics models, whose training losses Eq. (4) and (6) involve regressing on ∂ϕ ( s,a ) ∂s ˙ s and ∂ϕ ( s ′ ,a ) ∂s ′ ( -˙ s ) as target values. This results in a gradient through a gradient of ϕ ( · ) . Computationally, we calculate the Jacobian matrix ∂ϕ ( s,a ) ∂s using the vmap() function in Functorch † to ensure the second derivative of ϕ ( · ) can be correctly backpropagated during stochastic gradient descent. Similar a treatment can also be implemented with other auto-differentiation frameworks like Jax ‡ that support computing higher-order derivatives.

Pre-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 pre-training/training epochs for the experiments in this paper are reported in Table 2.

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 specific 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, we 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.

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 as the g ( z s + f ( z s , z a ) , z a ) produces similar values as -f ( z s , z a ) , resulting in small T-symmetry consistency loss values (Eq. (7)), however, the discrepancy between g ( z s ′ , z a ) and -f ( z s , z a ) 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:

Wefind 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

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 ∥ f ( z s , z a ) + g ( z s + f ( z s , z a ) , z a ) ∥ 2 2 and ∥ f ( z s , z a ) + g ( z s ′ , z a ) ∥ 2 2 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 3. Based on the different scales of the datasets, we basically only use two sets 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, we apply Dropout [56] with dropout rate of 0.1 to regularize the policy network in all tasks with 10k expert data.

B Detailed Experiment Setups

Reduced-size dataset generation. To create reasonable reduced-size D4RL datasets for a fair comparison, we sub-sample the trajectories in the datasets rather than directly sampling the ( s, a, s ′ , r ) transitions. For example, there are 2M ( s, a, s ′ , r ) transitions in the "halfcheetah-medium-expert" dataset, we first split these records into 2,000 trajectories based on the done condition, then randomly draw 10 trajectories (10k transition points) to serve as the reduced-size datasets for model training.

Experiment setups for representation learning evaluation. To evaluate the 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 as follows:

'AE-rep' model: We construct a vanilla auto-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 a latent forward dynamics prediction model f , which is implemented as a 4-layer feed-forward neural network with ReLU activation, and optimized using Adam optimizer (same as TDM). The forward model was trained by minimizing the loss term ∥ ˙ z s -f ( ϕ ( s, a )) ∥ 2 2 , where we directly regress f ( ϕ ( s, a )) with the ˙ z s derived from the latent states obtained from the encoder as ˙ z s = z s ′ -z s . Note that in this baseline, no ODE property nor T-symmetry regularization is included. Again we use the decoder to decode ˙ z s → ˙ s as 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 ∥ ˙ z s -f ( ϕ ( s, a )) ∥ 2 2 and ∥ ( -˙ z s ) -g ( ϕ ( s ′ , a )) ∥ 2 2 , where -˙ z s is directly calculated from the encoded latent states, i.e., ˙ z s = z s ′ -z s . Note that in this baseline, the T-symmetry is also implicitly captured, since both the latent forward and reverse dynamics models are regressing the same ˙ z s and its opposite value.

'SimSiam' model: For the self-supervised representation learning baseline, we implement an 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. [44].

Experiment setups for evaluating generalization performance. To evaluate TSRL's generalization capability beyond the offline datasets, we construct two low-speed datasets based on the original D4RL Walker2d medium and medium-expert datasets. In accordance with the Gym documentation, we selected the 'x-coordinate velocity of the top' (8th dimension of the states) in the walker environment to perform data filtering. We remove all samples with the x-coordinate velocity of the top greater than 0 . 2 × max-x-velocity 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 in Table 4. As shown in the results, TSRL achieves much better performance in the pen tasks, both the full datasets and the reduced-size datasets.

Additional results on Antmaze-umaze tasks. Wealso conduct experiments on the D4RL Antmazeumaze tasks with full and reduced-size 10k datasets. The results are presented in Table 5. We use the

same hyperparameters as in the D4RL MuJoCo 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 a recent model-based data augmentation method CABI [43]. CABI employs a pair of predictive dynamics models to assess the reliability of the augmented data. The results are presented in Table 6.

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.

Ablation on 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 if 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 ℓ fwd , ℓ rvs and ℓ T -sym 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 8. The experiment results demonstrated that an overly expressive model could not help the RL algorithm to derive a well-behaved policy with limited data

due to potential overfitting and inconsistency with the T-symmetry property. On the other hand, an overly regularized model may also hurt performance. This is consistent with our previous insight that a 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 in Figure 8. For each evaluation step, the policies are evaluated with 5 episodes over 3 random seeds.

Halfcheetah-expert-v2

Halfcheetah-medium-expert-v2

Halfcheetah-medium-replay-v2

Normalized Return

Normalized Return

40

Normalized Return

40

30

20-

10

101

8.0

8.00

8.0

0.25

0.50

0.75

1.00

0.25

0.50

0.75

1.00

0.25

0.50

0.75

1.00

Training Steps (x106)

Training Steps (x106)

Training Steps (x106)

Halfcheetah-medium-v2

Hopper-expert-v2

Hopper-medium-expert-v2

Normalized Return

Normalized Return

Normalized Return

5040

30-

20-

2

10

8.00

0.25

0.50

0.75

1.00

0.25

0.50

0.75

1.00

0.25

0.50

0.75

1.00

Training Steps (x106)

Training Steps (x106)

Training Steps (x106)

Hopper-medium-replay-v2

Hopper-medium-v2

Walker2d-expert-v2

40

Normalized Return

Normalized Return

30

20

10

.00

0.25

0.50

0.75

1.00

0.25

0.50

0.75

1.00

0.25

0.50

0.75

1.00

Training Steps (x106)

Training Steps (x106)

Training Steps (x106)

Walker2d-medium-expert-v2

Walker2d-medium-replay-v2

Walker2d-medium-v2

Normalized Return

Normalized Return

1

0.25

0.50

0.75

1.00

0.25

0.50

0.75

1.00

0.25

0.50

0.75

1.00

Training Steps (x106)

Training Steps (x106)

Training Steps (x106)

Peer reviews:
comments: 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.

STRENGTHS:
- The idea is interesting and makes intuitive sense.
- The paper is overall well-written.
- 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.
- No comparison of other offline reinforcement learning methods using data augmentation.
- No experiments to validate TSRL alleviates the problem of over-conservatism.

QUESTIONS:
- Explain the reason that the performance of TSRL is not comparable with other baselines in most Adroit human and cloned tasks.
- Add comparisons of other offline reinforcement learning methods using data augmentation.
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

comments: 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.

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.

2. The article is clearly structured.

3. The proposed method improves performance on offline RL with small datasets.

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.

2. The proposed method has a great improvement on small datasets, but there is still a lot of gap from the performance of complete data.

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?

2. 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?

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-05

comments: 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.

STRENGTHS:
- The proposed method makes sense in improving the sample efficiency of offline RL algorithms.
- 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.
- 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].
[1] Yu, Tao, et al. "Playvirtual: Augmenting cycle-consistent virtual trajectories for reinforcement learning." Advances in Neural Information Processing Systems 34 (2021): 5276-5289.

QUESTIONS:
1. Why can the proposed method improve the sample efficiency of off-line RL algorithms? Pls give in-depth and convincing analysis.
2. What are the relations between this work and those metioned in the weakness part?
3.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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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?", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-31"}