Title: Recovering from Out-of-sample States via Inverse Dynamics in Offline Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Offline reinforcement learning

state distributional shift

state recovery

inverse dynamics model
Abstract: In this paper we deal with the state distributional shift problem commonly encountered in offline reinforcement learning during test, where the agent tends to take unreliable actions at out-of-sample (unseen) states. Our idea is to encourage the agent to follow the so called state recovery principle when taking actions, i.e., besides long-term return, the immediate consequences of the current action should also be taken into account and those capable of recovering the state distribution of the behavior policy are preferred. For this purpose, an inverse dynamics model is learned and employed to guide the state recovery behavior of the new policy. Theoretically, we show that the proposed method helps aligning the transited state distribution of the new policy with the offline dataset at out-of-sample states, without the need of  explicitly predicting the transited state distribution, which is usually difficult in high-dimensional and complicated environments. The effectiveness and feasibility of the proposed method is demonstrated with the state-of-the-art performance on the general offline RL benchmarks.

Recovering from Out-of-sample States via Inverse Dynamics in Offline Reinforcement Learning

Abstract

Wedeal with the state distributional shift problem commonly encountered in offline reinforcement learning during test, where the agent tends to take unreliable actions at out-of-sample (unseen) states. Our idea is to encourage the agent to follow the so called state recovery principle when taking actions, i.e., besides long-term return, the immediate consequences of the current action should also be taken into account and those capable of recovering the state distribution of the behavior policy are preferred. For this purpose, an inverse dynamics model is learned and employed to guide the state recovery behavior of the new policy. Theoretically, we show that the proposed method helps aligning the transited state distribution of the new policy with the offline dataset at out-of-sample states, without the need of explicitly predicting the transited state distribution, which is usually difficult in high-dimensional and complicated environments. The effectiveness and feasibility of the proposed method is demonstrated with the state-of-the-art performance on the general offline RL benchmarks.

1 Introduction

Reinforcement learning has made significant advances in recent years, but it has to collect experience actively to gain understanding of the underlying environments. However, such online interaction is not always practical, due to either the potential high cost of data collection procedure or its possible dangerous consequences in applications as autonomous driving or healthcare. To address these issues, offline reinforcement learning aims to learn a policy from offline datasets without doing any actual interaction with the environments [19, 13, 1].

However, directly deploying online RL algorithms, such as Deep Deterministic Policy Gradient [18], to learn the new policy from the offline dataset without proper constraints would highly likely suffer from action distributional shift due to the change in the actions generated by the new policy. This would result in the so called extrapolation error [8], i.e., the TD target could be wrongly estimated when querying those out-of-distribution (OOD) actions generated by the new policy. To address these issues, methods like Conservative Q-Learning (CQL) [16], TD3+BC [7] and Implicit Q-Learning (IQL) [14], treat the OOD actions as conterfactual queries and try to avoid performing such queries completely during learning, hence suppressing the Q-function extrapolation error. However, such pessimism for out-of-sample data could be too restricted and sample inefficient, as not all out-ofsample(unseen) states are not generalizable [20].

To effectively generalize to out-of-sample or even OOD states, it is necessary to constrain the behavior of agent at those states, otherwise the policy extrapolation error that occurs at test stage may drive the

Behaviorpolicy

Tβ

(experttraj.)

CQLpolicy

πT

Traj.Divergence

knockedintoanout-of-samplestate

(Tβand

π

erroraccumulating

failure

timesteps

agent's transited state distribution away from the offline dataset, referred as the state distributional shift problem - a problem has not been paid enough attention as most previous works focus on addressing the issues of OOD actions during training [16, 7]. The problem of state distributional shift is likely to arise in those extensive and non-stationary environments, where it is impossible for the agent to cover the entire state space during training. A brief example is illustrated in Figure 1, where a robotic agent, trained using offline RL frameworks like CQL, is knocked by some external force unintentionally and falls into an out-of-sample state during test, which leads to failure finally due to the unreliable actions taken at this out-of-sample(unseen) state [3] and the error confounding effects.

To deal with the above problem, a natural idea is to teach the agent to recover from these out-ofsample states to their familiar regions, i.e., the demonstration of offline datasets. For this purpose, we propose the Out-of-sample State Recovery (OSR) method to implement this state recovery principle. Our idea is to estimate the actions whose consequence would be in the support region of the offline dataset using inverse dynamics model (IDM) [2]. The estimated IDM can be interpreted as an extended policy with pre-defined consequences, hence being suitable for guiding the learning of the new policy. Theoretically, we show that the proposed OSR makes the transited state distribution of the learnt policy at out-of-sample states align well with the offline dataset, without the need of building any forward dynamics models. We also propose a modified OSR-v that suppresses the probability of selecting actions according to the risk of OOD consequences through imposing extra value-constraints when decision making. We experimentally 2 demonstrated that the proposed OSR outperforms several closely related state-of-the-art methods in offline AntMaze/MuJoCo control suites with various settings.

In what follows, after an introduction and a review of related works, Section 3 provides a concise introduction to the background of offline RL and our work. The OSR method is presented in detail with theoretical analysis of its effectiveness in Section 4.2. The variant of OSR, called OSR-v, is then introduced with theoretical analysis of its advantages in Section 4.3. In Section 4.4, the practical implementation of both methods is described, including the loss function used. Experimental results are presented in Section 5 to evaluate the effectiveness of both methods under various settings. Finally, the paper concludes with a summary of the findings and contributions.

2 Related work

Robust offline reinfocement learning. Robustness is critical for offline RL methods to implement in real-world applications. Recent work RORL [25] adds smoothing term to relax the conservatism of algorithms like Conservative Q-Learning, making the agent generalize to out-of-sample states. ROMI [23] introduce reverse mechanism to build links between isolated data clusters based on out-of-sample data. However, it is still possible that the agent would drift away from the dataset as the results of the well-known error compounding effects due to the myopic of the new policy, referred as state distributional shift problem, which is more important in practical applications. For example, in healthcare [12], it is critical for the agent to make decisions with safe consequences, i.e., generating reliable trajectories. Theoretically, it is shown that control the divergence of the learnt policy from the behavior policy helps to bound the state distributional shift [22], but the bound is actually very loose [15]. Recently model-based State Deviation Correction (SDC) [27] builds a dynamics model and a transition model to guide the agent to generate in-distribution trajectories, through which to enhance the agent's safety and controllability. However, building a high-capacity dynamics model is not always practical in complicated applications, highlighting the necessity to constraint the agent's behavior without constructing any forward model to predict the high-dimensional observations.

Inverse dynamics model. An inverse dynamics model I ( a | s ′ , s ) predicts the action distribution that explains the transition between a given pair of states. Inverse dynamics models haven been applied to improving generalization to real-world problems [6], Markov representation learning [2], defining intrinsic rewards for exploration [5]. In our work, the inverse dynamics model is interpreted as a consequence-constraint policy that generates action distributions with predictable consequences, which in turn plays the role of supervision to the new policy.

3 Background

A reinforcement learning problem is usually modeled as a Markov Decision Process (MDP), which can be represented by a tuple of the form ( S, A, P, R, γ ) , where S is the state space, A is the action space, P is the transition probability matrix, R and γ are the reward function and the discount factor. Apolicy is defined as π : S → A and trained to maximize the expected cumulative discounted reward in the MDP:

In general, we define a Q-value function Q π ( s, a ) = E [ ∑ ∞ t =0 γR ( s t , π ( a t | s t )) | s, a ] to represent the expected cumulative rewards. Q-learning is a classic method that trains the Q-value function by minimizing the Bellman error over Q [24]. In the setting of continuous action space, Q-learning methods use exact or an approximate maximization scheme, such as CEM [11] to recover the greedy policy, as follows,

In offline setting, Q-Learning algorithms learns a Q-value function Q π ( s, a ) and a policy π from a dataset D , which is collected by a behavior policy π β . Since there is always an action distributional shift of the new policy π and the behavior policy π β , this basic recipe fails to estimate the Q-values for OOD state-action pairs. Conservative Q-Learning (CQL), as a representative OOD-constraint offline RL algorithm, tries to underestimate the Q-values for OOD state-action pairs to prevent the agent from extrapolation error [16]. The CQL term is as follows,

However, at test stage, there still exists a static state distributional shift of the new policy π and the behavior policy π β , which may accumulate the discrepancy between the agent and demonstration of the dataset. To migrate this, State Deviation Correction (SDC) [27] aims to train a policy choosing actions whose visited states are as closer to the dataset as possible. Specifically, the SDC trains a dynamics model M and a transition model U , to optimize the new policy, as follows,

where ˆ s is a perturbed version of the original state s , λ > 0 is the weight for the SDC regularization, and D is a distance measure, which is maximum mean discrepancy (MMD) in [27].

4 Out-of-sample state recovery using inverse dynamics

In this section, we first describe our Out-of-sample State Recovery (OSR) method in detail, as a policy regularization onto the actor loss, which is introduced in detail in section 4.2. Then we also propose a variant of OSR, OSR via value-constraints (OSR-v), for the purposes to reduce the hyperparameters and test the feasibility of practising state recovery principle in different modes, in section 4.3. Finally, we specifically describe how to implement OSR and OSR-v in practice in section 4.4.

4.1 Noise injection onto dataset

It is well-known that noise injection is helpful to address the covariate shift problem in deep learning [17]. While this technique is adopted in [27] to simulate the OOD states, in this work we interpret

the states perturbed with linear Gaussian noise as counter-examples on how to recover from those states into in-sample states. Given an offline dataset D , which consists of quadruples ( s, a, r, s ′ ) , we first perform data augmentation onto it to form a mixed dataset.

Specifically, given a quadruple ( s, a, r, s ′ ) , we perturb the state s to obtain a noisy state,

where ϵ is sampled from the standard Gaussian distribution N (0 , 1) and β is usually set to be a small constant. Besides, we could also utilize more complex methods to perturb the states, such as the adversarial attacks in [25]. In this way, we obtain a set of perturbed quadruples (ˆ s, a, r, s ′ ) , and group them into a new perturbed dataset ˜ D . Note that we do not perturb the next state s ′ , to preserve the destination we wish the agent to jump from ˆ s .

Finally we combine the two datasets D and ˜ D to form a new dataset D tot , i.e., D tot = D + ˜ D , where each element in D tot is denoted as (˜ s, a, r, s ′ ) . Next we describe how to train a inverse dynamics model to learn how to safely perform transition from { ˜ s } to s ′ .

4.2 Learning to recover from out-of-sample states via policy-constraints

A inverse dynamics model (IDM), denoted as I π β ( a | s, s ′ ) , is defined in terms of the behavior policy π β , dynamics model P ( s ′ | s, a ) and the transition function P ( s ′ | s, π β ) via Bayes' theorem, as follows,

where P ( s ′ | s, π β ) = ∑ a ∈ A P ( s ′ | s, a ) π β ( a | s ) . Note that to learn a IDM model, we don't have to estimate the dynamic model of the environment but only needs to treat this task as a usual function approximation problem using a neural network. In particular, using the quadruples (˜ s, a, r, s ′ ) in D tot , the needed IDM I π β ( a | s, s ′ ) can be estimated using a probabilistic regression model with ˜ s, s ′ as input and action a as output.

With the estimated IDM model available, we interpreted it as an extended policy with desired immediate consequence known, illustrating how to recover from out-of-samples into in-sample(safe) states, then use it to guide the learning of the new policy. In this work, we use the KullbackLeibler(KL) divergence to measure the divergence between two distributions. In particular, we minimize the KL divergence between I π β ( a | ˜ s, s ′ ) and the new policy π ( a | ˜ s ) , as follows,

where the ˜ s is sampled from the mixed dataset D tot and s ′ is its in-sample consecutive state. (7) is the OSR term aiming to play the same role of traditional model-based SDC, but without modeling the high-dimensional observations. Obviously, (7) utilizes a forward KL divergence, in which the difference between I π β ( a | ˜ s, s ′ ) and π ( a | ˜ s ) is weighted by I π β ( a | ˜ s, s ′ ) . Minimizing this forward KL divergence is equivalent to maximizing the policy likelihood as, max π E ˜ s ∼D tot ,s ′ ∼ P ( s ′ | s,π β ) ,a ∼ I π β ( a | ˜ s,s ′ ) [log π ( a | ˜ s )] , which is easy to achieve by the definition of KL divergence. In this manner, the learnt policy π is able to recover the average behavior of the IDM E s ′ ∼ P ( s ′ | s,π β ) I π β ( a | ˜ s, s ′ ) with greatest probability.

To further explain the feasibility of OSR term, Theorem 1 shows that the OSR term aligns the transited state distribution of the learnt policy with the offline dataset at out-of-sample states in the setting of bounded action space, which is enough to meet the requirements of most RL environments. First, two assumptions are given to serve Theorem 1 in Appendix A, where the assumptions give the continuity and positivity for the transition function of π β , and the positivity for the new policy we train. Then,

Theorem 1. Given two consecutive state s and s ′ , the out-of-sample state ˆ s within the ϵ -neighbourhood of s and the behavior policy π β . Given an inverse dynamics model I π β ( a | s, s ′ ) . In bounded action space setting, the following two optimization formulas are equivalent,

, that is, the two optimization formulas achieve the same policy.

The proof of Theorem 1 is conducted in Appendix B.1. It theoretically guarantees that OSR have the effect that the learnt policy is regularized to recover the transited state distribution P ( s ′ | ˆ s, π ) at out-of-sample state ˆ s aligned well with the transited state distribution P ( s ′ | s, π β ) of the behavior policy π β at the corresponding in-sample state s . For the fact that the offline dataset is generated via the behavior policy π β , then P ( s ′ | s, π β ) is the in-distributional transited state distribution of the offline dataset. Therefore, we remark that OSR helps aligning the new policy's transited state distribution P ( s ′ | ˆ s, π ) at out-of-sample state ˆ s with the offline dataset. Interestingly, we find that the term min π KL ( P ( s ′ | s, π β ) ∥ ∥ P ( s ′ | ˆ s, π ) ) could also be termed as a KL-divergence version of SDC introduced in (4).

4.3 Value-constraint out-of-sample state recovery: a variant

In this section, we propose Out-of-sample State Recovery via value-constraints (OSR-v) as a variant of OSR proposed before. The proposed OSR-v penalizes the Q-values of those actions that prefer transiting to OOD states at out-of-sample states, which drives the agent able to choose the actions that transits to in-sample states, recovering from out-of-sample states. In particular, given two consecutive states ˜ s and s ′ sampled from the mixed dataset D tot constructed before, we overestimate the Q-values of actions that transits to in-sample states, i.e., actions generated by the IDM I π β ( ·| ˜ s, s ′ ) , while suppressing the Q-values of actions with unknown consequences, i.e., actions generated by the new policy π , as follows,

where ˆ a tar refers to the action sampled from IDM, which is referred as target action.

In practice, we observe that OSR-v enable the agent to follow the state recovery principle as well. In order to rigorously explain this phenomenon, we give Theorem 2 under the assumptions mentioned in Appendix A to show that performing OSR-v is equivalent to underestimating the value of OOD states while overestimating the value of in-sample states, which guides the agent prefer to transit to those in-sample(safe) states, avoiding error accumulated via state distributional shift .

Theorem 2. Given two consecutive state s and s ′ , the out-of-sample state ˆ s within the ϵ -neighbourhood of s and the behavior policy π β . π is the new policy. Define Q ( s, a ) as a approximal Q-value function and V ( s ) as a approximal value function. Given an inverse dynamics model I π β ( a | s, s ′ ) . Then the following two optimization formulas are approximately equivalent:

, that is, they achieve the same policy approximately. Then the difference of the two optimization formulas is no larger than 2 δ ∑ a π β ( a | ˆ s ) Q (ˆ s, a ) .

The proof of Theorem 2 is conducted in Appendix B.4. Theorem 2 demonstrates the approximate equivalence of (9) and (10). There is actually a little gap between (9) and (10) which is caused by the existence of the continuous gap δ , so when δ is large, (9) and (10) are quite different. However, if we assume the infinite norm of reward function ∥ R ∥ ∞ ≤ 1 , then 2 δ ∑ a π β ( a | ˆ s ) Q (ˆ s, a ) ≤ 2 δ 1 -γ , where γ is the discount factor, and by the conservative selection of the magnitude of the noise ϵ , the δ would be a quite tiny value in practice.

The s ′ sampled from P ( s ′ | s, π β ) , the state distribution aligned with the offline dataset D , is in-sample, then overestimating of such s ′ increases the agent's preference to transiting to in-sample states, which is supported by the reasoning presented in Proposition 1, as follows,

Proposition 1. Given an arbitrary state s and its target state s ′ . Q is an approximal Q-value function and V is its value function. π Q is the learnt policy according to Q . Let π Q ( a | s ) is positive correlated to Q ( s, a ) , noted as π Q ( a | s ) ∝ Q ( s, a ) , and assume ∃ a, P ( s ′ | s, a ) > 0 , then P ( s ′ | s, π Q ) ∝ V ( s ′ ) .

The proof and more discussion of Proposition 1 is conducted in Appendix B.2. For the fact that out-of-sample(noisy) state ˆ s is close to s , the in-sample prior of s ′ , so we assume ∃ a, P ( s ′ | ˆ s, a ) > 0 , and via Proposition 1, E s ′ ∼ P ( s ′ | s,π β ) P ( s ′ | ˆ s, π ) ∝ E s ′ ∼ P ( s ′ | s,π β ) V ( s ′ ) . Therefore, the result of Theorem 2 and Proposition 1 guarantees that OSR-v can align the transited state distribution of the new policy with the offline dataset, avoiding compounding state distributional shift as well.

4.4 Implementation and algorithm summary

We implement our methods, OSR and OSR-v, introduced in previous sections based on an Conservative Q-Learning(CQL) framework to practice the state recovery principle in the setting of countinuous action space, without modeling the dynamics of the environment. To be specific, we would implement OSR in policy-constraint mode and OSR-v in value-constraint mode.

Policy-constraint mode (OSR). The policy-constraint based method, OSR proposed in (7), could be implemented as SR loss L sr in (11), where the detailed derivation is presented in Appendix B.5.

Then the problem is to estimate the KL divergence in (11) under the condition of continuous action space. In general, we often model the policy as π ( a | ˜ s ) = N ( µ π , σ π ; ˜ s ) and the IDM as I π β ( a | ˜ s, s ′ ) = N ( µ I , σ I ; ˜ s, s ′ ) , where N notes Gaussian distribution. First, we construct the SR loss for the policy network to estimate the KL divergence mentioned in (11).

There is an analytical solution for calculating the KL divergence between two Gaussian distributions, the Gaussian policy π ( a | ˜ s ) = N ( µ π , σ π ; ˜ s ) and the Gaussian IDM I π β ( a | ˜ s, s ′ ) = N ( µ I , σ I ; ˜ s, s ′ ) , so we can transfer (11) as the following SR loss:

where subscript d represents the value of d th dimension of the D -dimensional variable. After constructing the term of OSR in the settings of both discrete and continuous action space, we construct the actor loss function of the policy network as follows,

where λ is the balance-coefficient and the L sr is as (12) in the setting of continuous action space. L sr imply a behavior cloning from the target distribution estimated via the inverse dynamics model I π β ( ·| ˜ s, s ′ ) to our new policy to enable the agent always transits to those in-sample states s ′ from no matter in-sample s or out-of-sample ˆ s , i.e., ˜ s from the mixed dataset D tot mentioned in section 4.1. Then the critic loss function for the Q-networks is,

where α is the balance-coefficient; The second term is an one-step Bellman error, as (2) shows.

Value-constraint mode (OSR-v). We note the term of (8) mentioned in section 4.3, as L sr -v , to regularize the Q-networks. Then the critic loss function for the Q-networks is as follows,

where α is the balance-coefficient. We replace the CQL regularization, because L sr -v would degrade to CQL regularization when it performs on the in-sample part(the original offline dataset D ) of the mixed dataset D tot , as Proposition 2 shows.

Proposition 2. When training on the original dataset D , L sr -v is equivalent to CQL regularization.

The proof of Proposition 2 is conducted in Appendix B.3. Proposition 2 shows the feasibility of replacing the term of CQL by the OSR term, so that fewer hyperparameters are introduced. L sr -v overestimates the actions that prefer transiting to in-sample sates to avoid error accumulation. Then the actor loss L π we use in OSR-v is as follows,

To sum up, both OSR and OSR-v optimize actor loss L π to update the policy network π and critic loss L Q to update our Q-networks, and output the learnt policy network π . The whole process of OSR is summarized in Algorithm 1 in Appendix C while the whole process of OSR-v is summarized as Algorithm 2 in Appendix C.

5 Experiments

In experiments we aim to answer: 1) Does the proposed OSR help to achieve the state-of-the-art performance in offline RL? 2) Is the inverse dynamics model able to guide the policy recovering from out-of-sample states? 3) Is OSR able to improve the agent's robustness for out-of-sample situations? 4) Does OSR generalize well if the offline dataset only covers limited valuable area of the state space. Our experiments are organized as follows: Firstly, we conduct a comparative study on the MuJoCo and AntMaze benchmarks in the D4RL datasets for different versions of our method; then we design an out-of-sample MuJoCo setting to evaluate the generalization ability of CQL, SDC, the proposed OSR and OSR-v, and explore the behavior of the inverse dynamics model(IDM) and OSR policy; finally, we test the efficacy of OSR when presented with limited valuable data. A brief introduction of our code is available in Appendix D.9. MuJoCo. There are three types of high-dimensional control environments representing different robots in D4RL: Hopper, Halfcheetah and Walker2d, and five kinds of datasets: 'random', 'medium', 'medium-replay', 'medium-expert' and 'expert'. The 'random' is generated by a random policy and the 'medium' is collected by an early-stopped SAC [9] policy. The 'medium-replay' collects the data in the replay buffer of the 'medium' policy. The 'expert' is produced by a completely trained SAC. The 'medium-expert' is a mixture of 'medium' and 'expert' in half. AntMaze. AntMaze in D4RL is a maze-like environment where an ant agent must navigate through obstacles to reach a goal location. There are three different layouts of maze: 'umaze', 'medium' and 'large', and different dataset types: 'fixed', 'play' and 'diverse' which has different start and goal locations used to collect the dataset. Hyperparameter Details. We base the hyperparameters for OSR on those used in CQL [16]. We observe the performance of OSR is influenced by the noise magnitude β and SR weighting λ . Therefore, we conduct a sensitivity analysis of OSR in Appendix D.1.

Other experiments, including Sensitive Analysis, Ablation study, Comparison Study of implementations with Q-ensemble trick and so on, are attached in Appendix D.

5.1 Comparative study on offline MuJoCo/AntMaze suite

In this section, we compare our methods with several significant methods, including MOPO [26], RAMBO [21], IQL [14], CQL [16] and SDC [27], on the D4RL dataset in the MuJoCo benchmarks. Part of the results for the comparative methods are obtained by [27, 21]. All these methods are implemented without Q-ensemble trick [4] and results are shown in Table 1, where the proposed OSR and OSR-v achieve the similar or better performance than other methods on most tasks and the best average score, and we also remark that OSR and OSR-v generalize significantly better on the 'medium-expert' datasets where limited high-valuable data is available. More training details of OSR/OSR-v could be achieved in Appendix D.3. Besides, we also implement OSR with Q-ensemble trick of 10 Q-functions [4], termed as OSR-10, to compare with Q-ensemble based methods like RORL [25], SAC-10 and EDAC-10 [4], in Appendix D.4. From the results above, we conclude that OSR helps to achieve the state-of-the-art performance in offline RL.

5.2 Out-of-sample MuJoCo setting

To evaluate whether the inverse dynamics model and the proposed OSR policy enable the agent to recover from out-of-sample situations, we design Out-of-sample MuJoCo (OOSMuJoCo) benchmarks by adding different steps of random actions (Gaussian noise with magnitude of 1e-3) to simulate external force to knock the agent into out-of-sample states on Halfcheetah, Walker2d and Hopper. There are 3 levels of external force: slight(s) takes 5 steps of random actions, moderate(m) takes 10 steps and large(l) takes 20 steps. The visualization of OOSMuJoCo is shown in Figure 2.

Halfcheetah-OOS-

Hopper-OOS-

Hopper-OOS-

Hopper-OOS-

Walker2d-

Walker2d-OOS-

Walker2d-

OOS-slight

moderate

OOS-large

slight

moderate

large

OOS-slight

moderate

OOS-large

We evaluate the performance of policies trained by CQL, SDC, OSR and OSR-v on 'mediumexpert' datasets in the three benchmarks. The score and performance decrease of these policies across the 9 OOSMuJoCo benchmarks are shown in Table 2, where the performance decrease is measured as the percentage reduction in scores from OOSMuJoCo compared to the results of standard MuJoCo environments in Table 1. Our results indicate that the proposed OSR/OSR-v experiences significantly less performance degradation than the other two methods across most of the OOSMuJoCo benchmarks, demonstrating better robustness against perturbation for out-of-sample situations in the complicated and non-stationary real-world environments.

Besides, as a toy example, we visualize the Halfcheetah agent within an expert trajectory produced by the behavior policy π β on the offline dataset D , as is shown in Figure 3, then we also visualize the trajectories generated by the policy of OSR and the inverse dynamics model(IDM) guided from

the behavior policy, where we set the target state of IDM to the corresponding state in the expert trajectory. We add external force to knock the agent into an out-of-sample state and we observe that the agents equiped with IDM and OSR policy are both able to recover from the out-of-sample states. The state recovery processes of other 2 benchmarks is shown in Appendix D.5.2 and the visualizations of state distribution is shown in Appendix D.5.1. The results above demonstrate that the the inverse dynamics model is able to guide the policy recovering from out-of-sample state and OSR is able to improve the agent's robustness for out-of-sample situations.

BehaviorpolicyTβ

(expert traj.)

supervised

InversedynamicsI

三

Traj.Divergence

knockedintoanout-of-samplestate

recovering

recovered

（πβandI）

time steps

guidedviaIDM

OSRpolicy

π

Traj.Divergence

knockedintoanout-of-samplestate

recovering

recovered

（πβandπ）

timesteps

5.3 Limited valuable data setting

In this section, we test our method's feasibility when only limited valuable data is available for training, which is a novel task with the mixture of the expert dataset and random dataset with different ratios [27], aiming to evaluate the robustness of ORL algorithms on the condition that the offline dataset only covers limited valuable state space. In this paper, the proportions of random samples are 0.5, 0.6, 0.7, 0.8 and 0.9 and the ratios of expert samples decline, for Halfcheetah, Hopper and Walker2d. All datasets contain 1,000,000 samples.

Halfcheetah-random-expert

Hopper-random-expert

Walker2d-random-expert

100

80

40

60

20

88

67.8

88.5

44.5

90.9

90.2

46.7

87.8

85.2

42.6

86.9

88.4

34.7

70.2

79.5

115

105

100

110

112.4

112.8

111.2

111.0

111.6

111.1

112.4

110.4

112.0

110.8

112.8

112.3

108.5

102.2

111.1

150

100

50

88.8

111.9

88.3

65.9

106.7

81.5

110.2

48.3

85.5

49.8

106.5

50.5

106.7

26.8

0

95

0

17.6

0.5

0.6

0.7

0.8

0.9

0.5

0.6

0.7

0.8

0.9

0.5

0.6

0.7

0.8

0.9

random sample proportion

random sample proportion

random sample proportion

CQL

SDC

OSR(ours)

CQL

OSR(ours)

CQL

SDC

OSR(ours)

We compare the proposed OSR with CQL and SDC in a limited valuable data setting. The results, shown in Figure 4, demonstrate that all methods experience performance degradation when expert samples are limited. However, our proposed method, OSR, achieves the lowest performance decrease with the raise of the random sample ratio among the three methods in most tasks, which suggests that our method is less sensitive to the quality of training data. The results above mean that even if the dataset covers only limited valuable area of the extensive state space, the proposed OSR still works.

6 Conclusion

In this paper, we proposed a simple yet effective method, named Out-of-sample State Recovery (OSR), for offline reinforcement learning. The proposed method follows the state recovery principle to deal with the state distributional shift issue commonly encountered in offline RL. Unlike previous methods, we address this in a way that does not need to explicitly model the transition of the underlying environment by treating the inverse dynamics model as a guidance with desired immediate consequence for the new policy. Our theoretical results show that the proposed OSR is equivalent to aligning the transition distribution of the learnt policy with the offline dataset. The proposed method is evaluated on several offline AntMaze/MuJoCo settings and achieved the SOTA performance.

References

Mohammad Mehdi Afsar, Trafford Crump, and Behrouz H. Far. Reinforcement learning based recommender systems: A survey. CoRR , abs/2101.06286, 2021.

Cameron Allen, Neev Parikh, Omer Gottesman, and George Konidaris. Learning markov state abstractions for deep reinforcement learning. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 8229-8241, 2021.

Dario Amodei, Chris Olah, Jacob Steinhardt, Paul F. Christiano, John Schulman, and Dan Mané. Concrete problems in AI safety. CoRR , abs/1606.06565, 2016.

Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 7436-7447, 2021.

Jongwook Choi, Yijie Guo, Marcin Moczulski, Junhyuk Oh, Neal Wu, Mohammad Norouzi, and Honglak Lee. Contingency-aware exploration in reinforcement learning. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019 . OpenReview.net, 2019.

Paul F. Christiano, Zain Shah, Igor Mordatch, Jonas Schneider, Trevor Blackwell, Joshua Tobin, Pieter Abbeel, and Wojciech Zaremba. Transfer from simulation to real world through learning deep inverse dynamics model. CoRR , abs/1610.03518, 2016.

Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , pages 20132-20145, 2021.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA , volume 97 of Proceedings of Machine Learning Research , pages 2052-2062. PMLR, 2019.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018 , volume 80 of Proceedings of Machine Learning Research , pages 1856-1865. PMLR, 2018.

Geoffrey E. Hinton and Sam T. Roweis. Stochastic neighbor embedding. In Suzanna Becker, Sebastian Thrun, and Klaus Obermayer, editors, Advances in Neural Information Processing Systems 15 [Neural Information Processing Systems, NIPS 2002, December 9-14, 2002, Vancouver, British Columbia, Canada] , pages 833-840. MIT Press, 2002.

Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine. Scalable deep reinforcement learning for vision-based robotic manipulation. In 2nd Annual Conference on Robot Learning, CoRL 2018, Zürich, Switzerland, 29-31 October 2018, Proceedings , volume 87 of Proceedings of Machine Learning Research , pages 651-673. PMLR, 2018.

Hao-Cheng Kao, Kai-Fu Tang, and Edward Y. Chang. Context-aware symptom checking for disease diagnosis using hierarchical reinforcement learning. In Sheila A. McIlraith and Kilian Q. Weinberger, editors, Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18), the 30th innovative Applications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018 , pages 2305-2313. AAAI Press, 2018.

B. Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A. Al Sallab, Senthil Kumar Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey. IEEE Trans. Intell. Transp. Syst. , 23(6):4909-4926, 2022.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 . OpenReview.net, 2022.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing offpolicy q-learning via bootstrapping error reduction. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d'Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada , pages 11761-11771, 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

Michael Laskey, Jonathan Lee, Roy Fox, Anca D. Dragan, and Ken Goldberg. DART: noise injection for robust imitation learning. In 1st Annual Conference on Robot Learning, CoRL 2017, Mountain View, California, USA, November 13-15, 2017, Proceedings , volume 78 of Proceedings of Machine Learning Research , pages 143-156. PMLR, 2017.

Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Yoshua Bengio and Yann LeCun, editors, 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings , 2016.

Andrew Lobbezoo, Yanjun Qian, and Hyock-Ju Kwon. Reinforcement learning for pick and place operations in robotics: A survey. Robotics , 10(3):105, 2021.

Michael W. McCracken. Robust out-of-sample inference. Journal of Econometrics , 99(2):195223, 2000.

Marc Rigter, Bruno Lacerda, and Nick Hawes. RAMBO-RL: robust adversarial model-based offline reinforcement learning. CoRR , abs/2204.12581, 2022.

John Schulman, Sergey Levine, Pieter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. In Francis R. Bach and David M. Blei, editors, Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015 , volume 37 of JMLR Workshop and Conference Proceedings , pages 1889-1897. JMLR.org, 2015.

Jianhao Wang, Wenzhe Li, Haozhe Jiang, Guangxiang Zhu, Siyuan Li, and Chongjie Zhang. Offline reinforcement learning with reverse model-based imagination. arXiv e-prints , 2021.

Christopher J. C. H. Watkins and Peter Dayan. Technical note q-learning. Mach. Learn. , 8:279-292, 1992.

Rui Yang, Chenjia Bai, Xiaoteng Ma, Zhaoran Wang, Chongjie Zhang, and Lei Han. RORL: robust offline reinforcement learning via conservative smoothing. CoRR , abs/2206.02829, 2022.

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: model-based offline policy optimization. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

Hongchang Zhang, Jianzhun Shao, Yuhang Jiang, Shuncheng He, Guanwen Zhang, and Xiangyang Ji. State deviation correction for offline reinforcement learning. In Thirty-Sixth AAAI Conference on Artificial Intelligence, AAAI 2022 , pages 9022-9030. AAAI Press, 2022.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims to tackle a critical challenge in offline reinforcement learning, which involves recovering the state distribution during testing from out-of-sample states. To address this, the authors propose two methods, OSR and OSR-v, which leverage a learned inverse dynamics model to regularize the policy and underestimate the value function. The authors provide theoretical evidence supporting OSR's ability to align the transited state distribution of the new policy with the offline dataset distribution at out-of-sample states. Extensive experiments demonstrate that these proposed methods achieve state-of-the-art performance on general offline RL benchmarks as well as an out-of-sample MuJoCo benchmark.\n\nSTRENGTHS:\n* Overall, the paper is well-written and easy to follow.\n* This paper studies a significant problem in offline reinforcement learning, focusing on state distribution correction. The proposed methods show a link between State Deviation Correction (SDC) and robust offline RL.\n* The proposed method achieves state-of-the-art (SOTA) performance in both two general offline RL benchmark and an out-of-sample MuJoCo benchmark.\n* The effectiveness of OSR in recovering from specific perturbations is validated through visualizations presented in Section 5.2 and Appendix 4.2.\n* The code necessary for reproducing the results is provided.\n\nWEAKNESSES:\n* Although this paper has already conducted numerous experiments, I believe that a more comprehensive ablation study would be valuable. Specifically, in Appendix D.1, all the figures demonstrate that the smallest $\\beta$ achieves the best performance. I am curious to know the performance of OSR when $\\beta=0$, meaning the utilization of only $\\mathcal{D}$ for training. Theoretical expectations indicate that OSR-v's performance should align with CQL, while OSR's performance should match that of CQL+BC. If OSR with $\\beta=0$ still outperforms CQL and CQL+BC, it suggests that the action noise injected by the reverse model plays a crucial role in achieving a smooth policy output. Furthermore, conducting additional comparisons between CQL/CQL+BC with vanilla action noise would provide further insights.\n* I have concerns regarding Eq (5), as it assigns the same scale of Gaussian noise to different dimensions. Perhaps normalizing the states by their mean and standard deviation, similar to prior work [1], would be beneficial.\n* I think the proposed method also has a relevance to robust RL, as demonstrated by the Out-of-sample MuJoCo experiments which highlight OSR's capability to withstand testing-time attacks. Consequently, it is important to discuss prior works in robust RL and robust offline RL in the related works section, and to provide comparisons with existing robust offline RL approaches in the Out-of-sample MuJoCo benchmark.\n* The literature review lacks an inclusion of prior offline RL works such as ROMI [2], which incorporates a reverse model for data augmentation. Hence, the authors are suggested to further engage in a discussion and comparison with ROMI to enhance the comprehensiveness of the review.\n\n\n[1] Yang R, Bai C, Ma X, et al. Rorl: Robust offline reinforcement learning via conservative smoothing[J]. Advances in Neural Information Processing Systems, 2022, 35: 23851-23866.\n\n[2] Wang J, Li W, Jiang H, et al. Offline reinforcement learning with reverse model-based imagination. Advances in Neural Information Processing Systems, 2021, 34: 29420-29432.\n\nQUESTIONS:\n* Can the authors provide a more comprehensive ablation study, which should include scenarios such as OSR with $\\beta=0$, a comparison with CQL+BC, CQL+BC with state noise, and CQL+BC with both state and action noise? This will effectively illustrate the advantages of OSR and demonstrate the significance of incorporating state noise and the inverse model.\n* Would normalizing the states by their mean and standard deviation be more useful in addressing perturbations across different state dimensions and ensuring a more consistent $\\beta$?\n* Additionally, it is recommended to include an additional subsection in the related work that focuses on prior works in robust RL and robust offline RL.\n* It is suggested to compare prior robust offline RL work, such as RORL, against the out-of-sample MuJoCo benchmark\n* Furthermore, it is necessary to engage in further discussion and conduct comparisons with prior offline RL works that utilize a reverse model.\n\nLIMITATIONS:\nThe limitation highlighted by the authors in this work is the challenge of applying it in scenarios where the two basic assumptions, namely the smooth behavior policy transition and Gaussian policy assumption, do not hold.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-01"}

comments: SUMMARY:
The paper proposes a solution to the problem of state distributional shift in offline RL - the agent takes unreliable actions in out-of-sample states during testing. The paper introduces the use of inverse dynamics models to guide the state recovery behavior of learned policy. Without constructing forward model, OSR aligns the learned policy’s transited state distribution at out-of-sample state with the offline dataset. Experimental results demonstrate the effectiveness of the proposed method.

STRENGTHS:
I think the idea and effort to deal with state deviation in offline RL is meaningful and promising.

The paper is easy to read and understand.

The experimental results in the Out-of-sample MuJoCo setting are interesting and demonstrate the advantage of training state recovery policy.

WEAKNESSES:
It seems there is a non-negligible difference between the proposed method (theory) and the implementation. Eq. 7 and Eq.11 are not equivalent since the expectation w.r.t. s' is put outside the KL in Eq.11. One can find that they are not equal after simple calculation. I suppose this difference eliminates the need of predicting forward dynamics.

The paper lacks sufficient ablation study to support the effectiveness of the OSR/OSR-v term.

QUESTIONS:
How does OSR (not OSR-v) perform if the CQL term is removed? I think the OSR regularization term in Eq. 11 implies a kind of behavior cloning. Can this term alone suppress extrapolation error and overestimation? 

In Fig. 6, it seems that the weight $\lambda$ for the OSR regularization has little effect on the performance?

LIMITATIONS:
Although the performance in Tab. 1 is good, the hyperparamters need to be tuned per dataset (Tab. 3,4).
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the issue of state distributional shift in offline reinforcement learning, where an agent tends to take unreliable actions when faced with unseen states during testing. The authors propose a solution to encourage the agent to follow the state recovery principle when making decisions. In addition to considering long-term return, the agent takes into account the immediate consequences of its current action, prioritizing actions that are capable of recovering the state distribution of the behavior policy. To achieve this, the authors train an inverse dynamics model, which is then used to guide the state recovery behavior of the new policy. The authors demonstrate the effectiveness and feasibility of their approach by achieving state-of-the-art performance on general offline RL benchmarks. Importantly, the proposed method aligns the transited state distribution of the new policy with the offline dataset at out-of-sample states without the need for explicit prediction, which is particularly challenging in complex and high-dimensional environments.\n\nSTRENGTHS:\n1. The paper addresses the state distributional shift problem, which has been relatively overlooked in prior research that primarily concentrates on mitigating out-of-distribution (OOD) actions during training.\n\n2. The paper exhibits a high level of writing proficiency, logical organization, and reader-friendliness.\n\n3. The utilization of an inverse dynamics model as a guide for state recovery represents a novel approach in the field.\n\n4. The paper provides a rigorous theoretical analysis that demonstrates the effectiveness of the proposed algorithm.\n\nWEAKNESSES:\nA significant concern pertains to the construction of the noisy dataset. Two subproblems arise in this context. Firstly, the constructed noisy states may not adequately reflect the distribution of out-of-sample states encountered in the real world. For instance, when dealing with a robot control task, simply introducing Gaussian noise to the logged states might not align well with practical out-of-sample states due to various physical or environmental constraints on state transitions. Secondly, the constructed transition data, where the noisy state transitions to the next state, may not accurately represent the transition distribution in the real world. Consequently, this discrepancy between the two sources of data conflicts with the fundamental claim of the paper, as it becomes challenging to demonstrate that the inverse dynamics learned truly reflect the correct dynamics when the data sources are mismatched. More discussion or experimental results should be presented to address this concern.\n\nQUESTIONS:\nIn light of the aforementioned weaknesses, I have a few questions that, if addressed by the authors, could potentially enhance the overall quality of the paper:\n\n1. How can the issue of constructing a noisy dataset be mitigated to better reflect the distribution of out-of-sample states encountered in real-world scenarios, considering the presence of physical and environmental constraints on state transitions?\n2. Is there a way to ensure that the constructed transition data accurately represents the true transition distribution in real-world settings, thus aligning with the core claim of the paper regarding the learned inverse dynamics?\n3. Could the authors provide further evidence or explanations to alleviate concerns regarding the mismatch between the constructed noisy dataset and the actual out-of-sample states, thereby strengthening the validity of the proposed algorithm?\n\n\nI believe that addressing these questions would significantly contribute to improving the paper and potentially enhance its evaluation score.\n\nLIMITATIONS:\nOne notable limitation of this work is the potential mismatch between the constructed transition distribution and the true transition distribution encountered in real-world scenarios. This disparity may restrict the practical applicability of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors tackle the state distributional shift problem in offline reinforcement learning, by learning to *recover* to states that are close to the in-distribution region, where the proposed method is named Out-of-sample State Recovery (OSR). They augment the offline dataset by generating new samples with Gaussian noise injected to states for training. Also, they train an inverse dynamics model (IDM) to predict the *recovering* action given the current and next actions. Based on it, they learn a policy and encourage it to imitate the IDM to output *recovering* actions using their KL divergence term, or penalize Q-values for going out-of-sample. The authors present empirical results in MuJoCo and AntMaze environments.\n\nSTRENGTHS:\n- The training of IDM with perturbed samples and its use for encouraging the policy to recover to seen states is novel to some extent.\n- The experimental evaluation and analyses of the proposed methods are done in multiple settings and perspectives, which backs up the proposed methods empirically.\n- State distributional shift is one of the problems that need to be tackled in the field of offline reinforcement learning.\n\nWEAKNESSES:\n- The noise injection may not be enough in more complex environments with more state dimensions, as the state space would be too big to cover possible close OOD states with random sampling.\n- On a related note, as the size of the original offline dataset increases, the needed amount of augmented data could become too large or otherwise the augmentation might get less effective.\n- In terms of presentation, I think merging Fig.13 and Fig.14 would be more informative and make the comparison of the results with OSR and OSR-v easier.\n\nQUESTIONS:\n- Do you think the proposed IDM-based approach is supposed to work better than existing pessimism approaches with dataset augmentation (using noise injection) in general, especially when no specific modifications to cause out-of-sample situations are made for the environment? If so, I would like to hear the justification.\n\nLIMITATIONS:\nThe authors fairly covered the limitations of this work on the assumptions for the theoretical derivations and the recoverability depending on the type of out-of-sample situations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}