Recovering from Out-of-sample States via Inverse Dynamics in Offline Reinforcement Learning
Authors: Ke Jiang
Jia-Yu Yao
Xiaoyang Tan
Venue: NeurIPS 2023 poster (2023)
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
Ke Jiang! JJa-yu Yao Xiaoyang Tan 1.2+
1 College of Computer Science and Technology, Nanjing University of Aeronautics and Astronautics 2 MIIT Key Laboratory of Pattern Analysis and Machine Intelligence
School of Electronic and Computer Engineering, Peking University
ke ke_j jaajinngdn edu. jiayu_yao@pku. edu. edu. cn, x. tan@nuaa, edu. cn
Abstract
We deal with the state distributional shift problem commonly encountered in offline reinforcement learning during test, where the agent tends take unreliable actions at out-of out-of-sampll (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, inverse dynamics model is learned and employed to guide the state recovery behavior 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 -ft-ol-sample states, without the need of explicitly predicting the transited state distribution, which is usually difficult in high-ding digh-dimensional and complicated environments. The effectiveness and feasibility of the proposed method is demonstrated with the state-of- ooffthe-art performance on the general offline RI. benchmarks.
Introduction
Reinforcement learning has made significant advances in recent years, but it has collect experience actively to gain understanding of the underlying environments. However, such online interaction not always practical, due to either the potential high cost of data collection procedure or its possible dangerous consequences applications as autonomous driving or healthcare, To address these issues, offline reinforcement learning aims learn policy from offline datasets without doing any actual interaction with the environments 19
However, directly deploying online RL algorithms, such as Deep Deterministic Policy Gradient [88] 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 the so called extrapolation error [ i.c., the TD target could be wrongly estimated when querying those out-of- out-of-distribution (OOD) actions generated by the new policy. To address these issues, methods like Conservative Q-Leernning (CQL) 16). TD3+BC ( 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, not all out-of- sample(unseen) states are not generalizable [20]
To effectively generalize to oft-o--sample or even OOD states, it is necessary to constrain the behavior of agent those states, otherwise the policy extrapolation error that occurs at test stage may drive the
"Corresponding Author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Behavior policy (expert traj.)
COL policy x
Traj, Divergence "Rannd IT]
knocked into an out-oo--smppl state
time steps
Figure 1: An example of CQL agent's failure for the accumulative error due to state distributional shift Walker2d robotic agent. The interval between every two images is ten steps.
agent's '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 cover the entire state space during training. A brief example is illustrated in Figure 1 where a robotic agent, trained using offlin RL frameworks like CQL,, is knocked by some external force unintentionally and falls into out-of-sample state during test, which leads to failure finally due to the unreliable actions taken at this out-of oal-s--llmuuueeennn 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-of- sample states to their familiar regions, i..., the demonstration of offline datasets. For this purpose, we propose the Out-of 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) 22. 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 demonstrated a that the proposed OSR outperforms several closely related state-of state-of-the-ar methods in offline AntMaze/MuJJCo control suites with various settings.
In what follows, after an introduction and a review of related works, Section 3 provides a concise introduction the background of offline RL and our work. The OSR method 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 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, 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 healtheare [12, it is critical for the agent to make decisions with safe consequences, i.c., 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) (7 builds dynamics model and transition model to guide the agent to generate in-distribution trajectories, through which to enhance the agent's '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.
"ur code is available at huts://gthhubccom//kekk14453555

---

Inverse dynamics model. An inverse dynamics model [a]ssss) predicts the action distribution that explains the transition between 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 consequence-constrail policy that generates action distributions with predictable consequences, which in turn plays the role of supervision to the new policy.
Background
A reinforcement learning problem is usually modeled Markov Decision Process (MDP), which can be represented by a tuple of the form (S, A, P. R, 7),, where the state space, A is the action space, P is the transition probability matrix, R and are the reward function and the discount factor. policy is defined as т S + A and trained to maximize the expected cumulative discounted reward in the MDP:
max E
Σ
R((s T(ag()))]
In general, we define a Q-value function () RR(s [(aess))]ss, a] to represent the expected cumulative rewards. Q-learning 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 to recover the greedy policy, as follows,
+- arg R R(s. a) + 7B (16 -(,,,))))ss,)))"
π - arg max (. Q(s,a)
"
In offline setting, Q-Learning algorithms learns Q-value function (T a) and policy π from dataset D, which is collected by a behavior policy тt. Since there is always an action distributional shift of the new policy and the behavior policy *s, this basic recipe fails estimate the Q-values for OOD state-action pairs. Conservative Q-Learning (CQL), as representative OOD-constraint offline RL algorithm, tries underestimate the Q-values for OOD state-action pairs to prevent the agent from extrapolation error [16). The CQL term is as follows,
min E (Bs-Daasa Q(s. - Ee,a-D((sss)
(3)
However, at test stage, there still exists a static state distributional shift of the new policy т and the behavior policy Tg. 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 dynamics model M and a transition model U, to optimize the new policy, as follows,
closer to the dataset possible. Specifically, the SDC trains
min A DD (- (-15)),, U(:ss)
(4)
where & is perturbed version of the original state s, À > is the weight for the SDC regularization, and D is a distance measure, which is maximum mean discrepancy (MMD) in (27).
Out-of-sample state recovery using inverse dynamics
In this section, first describe our Out-of-sample State Recovery (OSR) method in detail, 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 implement OSR and OSR-v practice in section 4.4
4.1 Noise injection onto dataset
well-known that noise injection helpful address the covariate shift problem in deep learn- ing [17]. While this technique is adopted in [27] to simulate the OOD states, in this work we interpret
3

---

the states perturbed with linear Gaussian noise as counter- eeamteee 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 quadruple (8, a. r, s'), we perturb the state to obtain noisy state,
8 + 3
(5)
where is sampled from the standard Gaussian distribution N 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 (8, a, s') and group them into 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 8.
Finally we combine the two datasets and D form a new dataset Poots i.e. Prot = D, where each element in Dror is denoted as (3, o, 5'). Next we describe how to train a inverse dynamics model learn how safely perform transition from (3) to s'.
4.2 Learning recover from out-of-sample states via policy-constraint
inverse dynamics model (IDM), denoted as 1 (als, ). is defined in terms of the behavior policy Ty. dynamies model P(s"|s, a) and the transition function P(s [(s]s,, Tg) via Bayes' theorem, follows,
1 a s') (r((a)))P(5(ss
(6)
P(s'(s,73) m3)
where P()1(8;TT) = Σ P(s P(s')s,a)ms(ajp Note that to learn a IDM model, we don't have to aEA
estimate the dynamic model of the environment but only needs to treat this task as usual function approximation problem using a neural network. In particular, using the quadruples (3, a, s') in Diot, the needed IDM I (a|s, s') can be estimated using a probabilistic regression model with s, as input and action output.
With the estimated IDM model available, we interpreted it extended policy with desired immediate consequence known, illustrating how to recover from out-of-samples into in-sample(safe) states, then use it guide the learning of the new policy. In this work, we use the Kullback- Leibler(KL) divergence to measure the divergence between two distributions. In particular, we minimize the KL. divergence between 1Tt (a 8, s') and the new policy I [a as follows,
min Eg- KL ((g' (al5, ((a())
(7)
where the ă is sampled from the mixed dataset Diot and 8 is its in-sample consecutive state. (7) is the OSR term aiming to play the same role of traditional model-based SDC, but with- out modeling the nigh-dimensional observations. Obviously, (7) utilizes a forward KL diver- gence, in which the difference between - (a|5,, s') and (a is weighted by 1 (a s'). Minimizing this forward KL divergence is equivalent to maximizing the policy likelihood as, max, S' plel""""eee...).aa>=="""" (alá,a"" [log T (a[8)], which is easy to achieve by the definition of KI. divergence. In this manner, the learnt policy ι is able to recover the average behavior of the IDM P(gl-P(a 7 (a88, 8') with greatest probability.
To further explain the feasibility of OSR term, Theorem Saarrem[shore shows that the OSR term aligns the transited state distribution of the learnt policy with the offline dataset at oft-ol-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 I in Appendix A where the assumptions give the continuity and positivity for the transition function of Tg. and the positivity for the new policy train. Then,
Theorem 1. Given two consecutive state and s', the out- out-of-sample state å within the c- neighbourhood of a and the behavior policy ng. Given an inverse dynamics model 1 (als, 8'). In bounded action space setting, the following two optimization, formulus are equivalent.
min K L. (E (Bgg T (al8,s') r(als)) - + min MiinKK( (P(dlp T3) (( (())((.
that is, the two optimization formulas achieve the same policy.

---

The proof of Theorem is conducted in Appendix B.I It theoretically guarantees that OSR have the effect that the learnt policy is regularized recover the transited state distribution ((s||s, x) at out-of-sample state aligned well with the transited state distribution P(s'|s, π&) of the behavior policy ag at the corresponding in-sample state 8. For the fact that the offline dataset is generated via the behavior policy R3, then P(s'|s, (s, TT) 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'|3, 1) at out-of-sample state å with the offline dataset. Interestingly, we find that the term min, K ( ((((')ss (T))|| PPss) ()) 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 in-sample states, recovering from out-of-sample states. In particular, given states & and s' sampled from the mixed dataset Dint constructed before, we overestimate the Q-values of actions that transits in-sample states, i.e., actions generated by the IDM 1 (-18, 8'), while suppressing the Q-values of actions with unknown consequences, i.c., actions generated by the new policy æ, as follows,
consecutive
min Di (E)) ((5, -a (113,,)) (5, asar)
where åear refers to the action sampled from IDM, which is referred target action.
In practice, we observe that OSR-v enable the agent to follow the state recovery principle as well. In order rigorously explain this phenomenon, we give Theorem@undee the assumptions mentioned in Appendix AA 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 and s', the out-of-sample state å within the neighbourhood of s and the behavior policy πg. is the new policy. Define ((s a) as approximal Q-value function and Vss) as a approximal value function. Given an inverse dynamics model Ihe (als, Then the following two optimization formuul
formulas are approximately equivalent:
min E (18)) Q(s, à) - Eaia Ba (118,')) Q($, aear)
(9)
& min E. P(er V(d) BeleP(e))s.rr)) V(5' V
(10)
that is, they achieve the same policy approximately. Then the difference of the two optimization formulas is no larger than 26 >n, [n(als)Q($.aa)
The proof of Theorem 2 is conducted 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 6, when is large, g and (T0) are quite different. However,
if we assume the infinite norm of reward function| R| I then 26 Σ. (a(a a) < 26 T-
where y the discount factor, and by the conservative selection of the magnitude of the noise e, the would be a quite tiny value in practice.
The s' sampled from P(s'|s, T.). the state distribution aligned with the offline dataset D, is in in-sample, then overestimating of such s' increases the agent's 's preference to transiting to in in-sample states, which is supported by the reasoning presented in Proposition ! as follows,
Proposition 1. Civen an arbitrary state s and its target state s'. Q is an approximal Q-value function and V is its value function. TQ is the learnt policy according Q. Let TQ (a(a)s) s) is positive correlated to Q(s, a), noted as TQ(aas x Q(s, a), and assume 3a, P(s'|s., a) > 0, then P(s'|s, TQ) x (s')..
The proof and more discussion of Proposition I is conducted in Appendix B.2 out-o--sample(noiss) state is close to 8, the in-sample prior of s', so we assume Da, P(s'I8,a) > and via Proposition [] (e/re(ee))) P(s'|è (8, T) E.g'>(a"'l..rg) (s'). Therefore, the result of Theorem 2 and Proposition]] guaranteess that OSR-v can align the transited state distribution of the new policy with the offline dataset, avoiding compounding state distributional shift well.
For the fact that
5

---

4.4 Implementation and algorithm summary
We implement our methods, OSR and OSR-v, introduced in previous sections based on an Conserva- tive Q-Lea)ning(CLL) 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-comm colicy-constraint based method, OSR proposed in 7 could be implemented as SR loss Ler in (1). where the detailed derivation is presented in Appendix B.5
Lgr Da KL T (al5, 8') )
(11)
Then the problem is to estimate the KL divergence in (TI under the condition of continuous action space. In general, we often model the policy ((a (N Fxi3) and the IDM as I (a]s, s') = N (pz,
01; 01 s'), where /V' notes Gaussian distribution. First, we construct the SR loss for the policy network to estimate the KL divergence mentioned in (TT).
There is an analytical solution for calculating the KL divergence between two Gaussian distributions, the Gaussian policy (aa() = =N(u,, Tw; 8) and the Gaussian IDM I (alā, 8') (ur, 8'), so we can transfer as the following SR loss:
D
of (of)² - (uf u )² 2(
Lgr Ea
E(og d=1
(12)
o
where subscript d represents the value of dh dimension of the D - dimensiooal variable. After constructing the term of OSR in the settings of both discrete and continuous action space, construct the actor loss function of the policy network as follows,
Ls Egn D Q(5, a) + AL
where À is the balance-coefficient and the Ler is 12) in the setting of continuous action space. Lar imply a behavior cloning from the target distribution estimated via the inverse dynamics model (-18, 8") to our new policy to enable the agent always transits to those in-sample states s' from no matter samsample s or out-of-sample 5, i.c., from the mixed dataset Dot mentioned in section 4.1 Then the critic loss function for the Q-networks is,
LQ = E, a (Ea- ((s, a) Q(s,a)) + (r + Q(s', ((s',a') - Q(s,((a)) "
where a is the balance-coefficientt The second term is an one-step Bellman error, as (2) shows.
Value-constraint mode (OSR-v). We note the term of S mentioned in section 4.3 as Lar to regularize the Q-networks. Then the critic loss function for the Q-networks is as follows,
LQ a Lor-v + E.a.ree [r - Q('s a') - Q(s,a))
(13))
where Tsr-o= Ba,o'
E [aa ((,
B Q( âtar)
where a is the balance-coefficient. We replace the CQL regularization, because Ler-n would degrade to CQL regularization when it performs on the in-sample part(the original offline dataset D) of the mixed dataset Diot. as Proposition[] shows.
Proposition 2. When training the original dataset DD Ler-y is equivalent to CQL regularization.
The proof of Proposition 2 is conducted in Appendix B.3 replacing the term CQL by the OSR term, so that fewer hyperparameters are introduced. L overestimates the actions that prefer transiting to in-sample sates to avoid error accumulation. Then the actor loss Li we use in OSR-v is as follows,
Proposition 2 shows the feasibility of
L Eg D [Q(5,a)
To sum up, both OSR and OSR-v optimize actor loss L, to update the policy network π and critic loss Lo to update our Q-networks, and output the learnt policy network æ. The whole process of OSR summarized in Algorithm in Appendix C while the whole process of OSR-v is summarized
Algorithm 2 in Appendix

---

Table 1: Results of OSR(ours), OSR-v(ours), SDC, RAMBO, MOPO, IQL and CQL on offiin MuloCo and AntMaze control tasks averaged over seeds. indicates the average without 'expert' datasets. The top-2 highest scores in each benchmark are bolded.
Task name
CQL
IQL
MОPO
RAMBO
SDC
OSR
OSR-v
Halfcheetah-r.
35,4
13.1
31.9
40.0
36.2
35.2 1.2
35.9 1.0
Walker2d-r.
7.0
5.4
13.3
11.5
14.3
13.5 2.8
14.1 3,9
Hopper-r.
10.8
7.9
13.0
21.6
10.6
10.3# 3.1
13.53 3.4
Halfcheetah-m.
44.4
47.4
40.2
77.6
47.1
48.81 0.2
49.1
Walker2d-m.
79.2
78.3
26.5
86.9
81.1
85.7=
85.1 0.7
Hopper-m.
58.0
66.2
14.0
92.8
91.3
83.1 1.9
95.2:: 1.1
Halfcheetah-m.-rt
46.2
44.2
54.0
68.9
47.3
46.81 0.4
48.5: 0.6
Walker2d-m.rr.
26.7
73.8
92.5
85.0
30.3
87.9 1.1
87.8 1.0
Hopper-m.-t
48.6
94.7
42.7
96.6
48.2
96.71 1.7
101.2 0.8
Halfcheetah-m.-e
62.4
86.7
57.9
93.7
101.3
94.7+ 0.9
99.1 0.4
Walker2d-m.-e.
98.7
109.6
51.7
68.3
105.3
114.3+ 1.5
112.9 0.8
Hopper-m.-e.
111.0
91.5
55.0
83.3
112.9
113.1 0.6
113.2 0.6
Halfcheetah-e.
104.8
95.0
106.6
9.7
102.1 0.8
Walker2d-e.
153.9
109.4
108.3
110.3 0.4
111.4 0.8
Hopper-e.
109.9
109.9
112.6
113.1 0.5
112.9 0.4
MuJoCo-v2 Avg.
66.5
68.9
41.1*
68.5*
70.2
76.7(69.2
78.9(71.5*)
AntMaze-u.
74.0
87.5
0.0
25.0
89.0
89.9
92.0-10.8
AntMaze-m.-pp
61.2
71.2
0.0
16.4
71.9
66.0+2.5
11.
AntMaze-l.-p.
15.8
39.6
0.0
0.0
37.2
37.9+2.7
38.
AntMaze-u.-d.
84.0
62.2
0.0
0.0
57.3
74.042.8
69.0
AntMaze-m.-d.
53.7
70.0
0.0
23.2
78.7
80.0+2.7
77.0+2.44
AntMaze-l.-d.
14.9
47.5
0.0
2.4
33.2
37.9+1.5
34.0+2.22
AntMaze-vO Avg.
50.6
63.0
0.0
22.3
61.2
64.3
63.4
5 Experiments
In experiments aim to answer: 1) Does the proposed OSR help to achieve the oftte-of-theaat 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 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 follows: Firstly, we conduct comparative study on the MuJoCo and AntMaze benchmarks in the D4RL datasets for diffferent versions of our method; then we design an out-ol-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. brief introduction of our code is available in Appendix D.9 Mu.ooC. 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', mmedium-repaa',, 'medium-expee The 'random' 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 meddim-exppert and 'eppert' in half. AntMaze. AntMaze in D4RL 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 those used in COL [16). We observe the performance of OSR is influenced by the noise magnitude 33 and SR weighting Therefore, we conduct a sensitivity analysis of OSR in Appendix D.1
and 'expert'.
a mixture 'medium' agent
maze- like environment where
Other experiments, including Sensitive Analysis, Ablation study, Comparison Study of implementa- tions with Q -eesemblle trick and attached in Appendix

---

5.1 Comparative study on offline MuJoCo/AntMaaze suite
In this section, we compare our methods with several significant methods, including MOPO (66). RAMBO [21). IQL [4). 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. these methods are implemented without Qeensemble trick [4] and results are shown in Table I 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 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, also implement OSR with Qensemble trick of 10 O-functions 4 termed as OSR-10, to compare with Q-ensemble based methods like RORL [25]. SAC-10 and EDAC-10 A Appendix [.4] From the results above, 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 inverse dynamics model and the proposed OSR policy enable the agent to recover from out-of-sample situations, we design Out-of-sample MuJoCo (OOSMuloCo) benchmarks by adding different steps random actions (Gaussian noise with magnitude of Ic-3) to simulate external force to knock the agent into out-of-sample states Halfcheetah, Walker2d and Hopper. There are 3 levels external force: slight(s) takes steps of random actions, moderate(m) takes 10 steps and large(I) takes 20 steps. The visualization of OOSMuJoCo is shown in Figure 2
Halkhertdd CO5-light
HJK War-wrah-COy DUs
IaFchwaat- DO5 lage
llpppenOOs-
Pappse COS
Viv bwit- Dos-llght
Weleet) OUS
Wallard- arge
node sit
right
brga
Figure The visualizations of the 9 OOSMuJoCo benchmarks. The red arrows represent the direction the agents move after external force, and the agents fall into out-of-sample states then.
We evaluate the performance of policies trained by CQL,, SDC, OSR and OSR-v on 'medium- expert' 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 measured the percentage reduction in scores from OOSMuJoCo compared 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 the OOSMuJoCo benchmarks, demonstrating better robustness against perturbation for out-of -ft-ol-sapple situations in the complicated and -tatiootonnay real- world environments.
Table Results of CQL., SDC, OSR/OSR-v OOSMuloCo setting on the normalized return and decrease metric averaged over 4 seeds. The hightest score and lowest decrease are bolded.
CQL score
SDC score
OSR
OSR-v score
Task name
dec.(%)
dec. (%)
score
dec
dec.
Halfcheetah-OOS-..
56.2
10.0
100.9
0.4
94.1
0.6
98.4
0.7
Halfcheetah-OOOS..
52.9
15.3
99.8
1.5
92.7
2.1
98.4
0.7
Halfcheetah-OOSS-L
42.1
32.5
83.3
17.8
91.7
3.2
79.8
19.5
Walker2i-OOS-s.
81.1
47.3
102.7
2.5
113.3
0.9
110.3
2.3
Walker2d-00S
77.4
21.6
102.2
2.9
113.2
1.0
110.9
1.8
52.0
47.3
95.6
9.2
110.1
3.7
106.7
5.5
HOpper-OOS-s.
110.4
0.6
112.0
0.6
111.4
1.5
112.5
0.6
Hopper-OOS-m.
97.3
12.3
110.8
1.9
109.2
3.4
112.2
0.9
HOpper-OOS-I.
54.4
51.0
89.2
20.9
106.1
6.2
105.6
6.7
Besides, toy example, we visualize the Halfcheetah agent within an expert trajectory produced by the behavior policy ny on the offline dataset DD as is shown Figure 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 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 benchmarks is shown in Appendix [D.5.2 and the visualizations of state distribution is shown in Appendix D 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.
Behovior policy Hy (expert traj] supervised
dynamics
Divergence
knocked saa--a--aaple
45 and
time steps
guided via IDM
OSR policy
Traj. Diverpence agannd #1
knocked into an out sample state
recovering
steps
Figure 3: The state recovery process via IDM and OSR on the Hallichetah--OS--age benchmark. The interval between every two images ten steps.
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 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, 1,000,000 samples.
Hallihoetahh randoms oxpert
Hogper-random rondom expert
Waller2di cww.o/dddanddmmmooommm oapert
$05
X5
1120
COL
SDC
ame SDC
a (SR))ou)))
SDC
Figure 4: The curve of CQL, SDC and OSR in the setting of limited valuable data.
We compare the proposed OSR with CQL and SDC in a limited valuable data setting. The results, shown in Figure 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 the random sample ratio among the three methods in most tasks, which suggests that 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.
Conclusion
In this paper, we proposed simple yet effective method, named Out-of-sample State Recovery (OSR), for offline reinforcement learning. The proposed method follows the state recovery principle 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 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/MullCo settings and achieved the SOTA performance.

---

Acknowledgement
This work is partially supported by National Key R&D program of China (2021ZD01 13203), National Science Foundation of China (51976115).
References
[1] Mohammad Mehdi Afsar, Trafford Crump, and Behrouz. H. Far. Reinforcement learning based recommender systems: A survey. CoRR, abs/2101 10.20666666, 2021.
[2] Cameron Allen, Neev Parikh, Omer Gottesman, and George Konidaris. Learning markov state abstractions for deep reinforcement learning. In Mare' 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, NewrlPS 2021, December 6- 14, 2021, virtual, pages 9229-82211 2021.
[3] Dario Amodei, Chris Olah, Jacob Steinhardr, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in Al safety. CoRR, wbs///06..eeeemm 66560, 2016.
[4] Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty- offline reinforcement learning with diversified q-enseblee. 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.
[5] Jongwook Choi, Yijie Guo, Marcin Moczulski, Junhyuk Oh, Neal Wu, Mohammad Norouzi, and Honglak Lee. -ontingency-aware exploration in reinforcement learning. In 7th International Conference Learning Representations, ICIR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[6] Paul F, Christiano, Zain Shah, Igor Mordatch, Jonas Schneider, Trevor Blackwell, Joshua Tobin, Pieter Abbeel, and Wojeiech Zaremba. Transfer from simulation real world through learning deep inverse dynamics model. CoRR, abs/1610.03318, 2016.
[7] Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning, In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wort- man 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-20144 2021.
[8] 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.
[9] Tuomas Haamoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critie: Off- 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 Proceddings of Machine Learning Research, pages 1856-1865. PMLR, 2018
[10] 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 933-840. MIT Press, 2002.
[11] Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine. Seal- able 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-67 PMLR, 2018
10

---

[12] Hao-Cheng Kao, Kai-Fu Tang, and Edward Y. Chang. Context-aware symptom checking for disease diagnosis using hierarchical reinforcement learning. Sheila A. Mcllraith and Kilian Q. Weinberger, editors, Proceedings of the Phirty-Second AAAI Conference on Artificial Intelligence, (AAAI-18) 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
[13] B. Ravi Kiran, Ibrahim Sobh, Victor Talpaert, Patrick Mannion, Ahmad A. Al Sallab, Senthil Ku- mar Yogamani, and Patrick Pérez. Deep reinforcement learning for autonomous driving: A survey, IEEE Trans. Intell. Trunsp. Syst., 23(6):4909-4926, 2022.
[14] lya 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.
[15] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learming via bootstrapping eπror reduction. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buss Emily B. Fox, and Roman Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 11761-11771, 2019.
[16] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q 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 12, 2020, virtual, 2020.
[17] 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.
[18] 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.
[19] Andrew Lobbezoo, Yanjun Qian, and Hyock-Ju Kwon. Reinforcement learning for pick and place operations in roboties: A survey. Robotics, 10(3):105, 2021.
[20] Michael W. McCracken. Robust out-of-sample inference. Journal of Econometrics, 99(2):195- 223, 2000.
[21] Mare Rigter, Bruno Lacerda, and Nick Hawes. RAMBO-RL: robust adversarial model-based offline reinforcement learning. CoRR, abs/2204. 125122204.2211, 2022.
[22] 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.
[23] 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.
[24] Christopher J. C. H. Watkins and Peter Dayan. Technical note q-learning. Mach. Learn., 8:279-292, 1992.
[25] Rui Yang, Chenjia Bai, Xiaoteng Ma, Zhaoran Wang, Chongjie Zhang, and Lei Han. RORL: robust offline reinforcement learning via conservative smoothing. CoRR, abx/2206.0282222 2022.
11

---

[26] 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, Ad- vances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6- 12, 2020, virtual, 2020.
[27] Hongchang Zhang, Jianzhun Shao, Yuhang Jiang. Shuncheng He, Guanwen Zhang, and Xi- angyang Ji. State deviation correction for offline reinforcement learning. Thirty-Sixii Sixth AAAI Conference on Artificial Intelligence, AAAI 2022, pages 9022-9033 AAAI Press, 2022.
12


Reviews:
[{"review_id": "oVWoC1LKhP", "reviewer": "Reviewer_tZMd", "summary": "The 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.", "questions": "- 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.", "limitations": "The 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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thoughtful comments. We provide clarification to your questions as below. We appreciate it if you have any further questions or comments.\n\n**W1: The noise injection ... not be enough in more complex environments with more state dimensions...**\n\n**Response**: We agree that Gaussian noise injection would quickly become challenging in modeling OOD in high dimensional state space. In our opinion, an ideal OOD sample generator should be both efficient  or purposeful and relatively insensitive to the dimensions of the state space. One possible way to achieve this is through adversarial attacks [1] which purposefully searches those OOD samples in each $s$ centered $\\epsilon$-neighborhood such that they either cause large policy perturbation or have lower Q value. The method is also less sensitive to the dimensions of state space due to the technique of deep generative networks. We construct a new dataset using this method and denote the experimental results on it as OSR-a, shown below,\n\n| |Halfcheetah-m.|Halfcheetah-m.-r.|Halfcheetah-m.-e.|Halfcheetah-r.|Halfcheetah-e.|Hopper-m.|Hopper-m.-r.|Hopper-m.-e.|Hopper-r.|Hopper-e.|Walker2d-m.|Walker2d-m.-r.|Walker2d-m.-e.|Walker2d-r.|Walker2d-e.|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|OSR-a|**52.3**|**51.2**|**101.2**|**35.2**|**103.1**|**85.2**|**97.4**|112.7|**10.9**|112.2|**87.2**|86.8|113.4|**15.1**|**111.2**|\n|OSR|48.8|46.8|94.7|**35.2**|97.7|83.1|96.7|**113.1**|10.3|**113.1**|85.7|**87.9**|**114.3**|13.5|110.3|\n\nFrom the results we observe that OSR-a outperforms OSR at most benchmarks. We remark that better-designed out-of-sample generating methods would further improve the performance of our method, probably because such methods take into account the physical and environmental constraints, resulting in better sampling efficiency. \n\n**W2: ... dataset increases, the needed amount of augmented ...**\n\n**Reponse:** Please note that modeling OOD samples based on the observed experience is not our ultimate goal. Actually, as mentioned in our response to Q1 of Reviewer cN6c, the performance of our method is largely due to the inverse model, although noise injection is shown to further improve the performance - for your convenience, we copy the results below, where our method of OSR($\\beta=0$) is roughly equivalent to CQL + IDM without noise injection, while OSR is with noise injection. From these we can see that it is IDM, instead of noise injection, that contributes most to the performance of OSR, and data augmentation is more like icing on the cake for OSR (see 2nd and 3rd rows). Hence we conclude that even with no noise injection at all, our method would preserve most of its advantages in training a rational agent in the offline RL setting, provided that the dataset is large enough to train a good inverse dynamics model (IDM) . \n\n| |Halfcheetah-m.-e.|Hopper-m.-e.|Walker2d-m.-e.|\n|-|-|-|-|\n|CQL|62.4|98.7|111.0|\n|OSR($\\beta$=0)|92.3|111.8|110.1|\n|OSR|**94.7**|**114.3**|**113.1**|\n\n**W3: ... merging Fig.13 and Fig.14 ...**\n\n**Response:** Thanks for your suggestion and we will do that in the revised manuscript.\n\n**Q1: .... OSR ... better than existing pessimism approaches with dataset augmentation (using noise injection) .... justification.**\n\n**Response:**  Our answer is yes, and our justification is as follows, \n\nAlthough data augmentation is useful when the training data are not so representative, encountering out-of-sample situations is almost inevitable for any agent due to the presence of approximal errors and other factors.  Such unfortunate situations aggravate when the new policy makes unreliable decisions at unseen states, causing the agent to deviate from the offline dataset, and existing data augmentation-based approaches, e.g., RORL [1] and ROMI [2], seldom consider the problem of how to recover from such scenarios. \n\nTo give some experimental evidence on this, we compare two popular data augmentation approaches, i.e., ROMI and RORL,  with our method on the standard MuJoCo benchmark in the table below, where OSR-10 is OSR implemented with 10 Q networks,  the same as RORL. \n\n| |Halfcheetah-m.|Halfcheetah-m.-r.|Halfcheetah-m.-e.|Halfcheetah-r.|Hopper-m.|Hopper-m.-r.|Hopper-m.-e.|Hopper-r.|Walker2d-m.|Walker2d-m.-r.|Walker2d-m.-e.|Walker2d-r.|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|ROMI|49.1|47.0|86.8|24.5|72.3|98.1|111.4|30.2|84.3|**109.7**|109.7|7.5|\n|RORL|66.8|61.9|107.8|**28.5**|104.8|102.8|112.7|**31.4**|**102.4**|90.4|121.2|21.4|\n|OSR(ours)|48.8|46.8|94.7|35.2|83.1|96.7|113.1|10.3|85.7|87.9|114.3|13.5|\n|OSR-10(ours)|**67.1**|**64.7**|**108.7**|28.1|**105.5**|**103.1**|**113.2**|30.2|102.0|93.8|**123.4**|**23.1**|\n\nThese results show that although only a very naive noise injection mechanism is used for data augmentation in our OSR, its performs comparable or better than both ROMI and RORL in which more complex data augmentation techniques are involved. Besides, we also compare our method with RORL in the Out-of-sample MuJoCo benchmarks, and the results could be found in our response to the reviewer cN6c, showing that our IDM-based offline RL method OSR works much more stable than the state of the art data augmentation-based method (RORL) on the challenging OOD benchmarks. \n\nFurthermore, in the attached PDF file, we illustrate visually how our method enables the agent to recover from many out-of-sample situations on MuJoCo benchmarks.\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[J].arXiv e-prints, 2021.DOI:10.48550/arXiv.2110.00188."}, {"review_id": "ZFB0LFx6OZ", "reviewer": "Reviewer_tm1v", "summary": "The 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.", "questions": "In 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.", "limitations": "One 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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "A 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comment. We appreciate your questions and provide clarification below. \n\n**Q1: ... noisy dataset be mitigated to better reflect the distribution of out-of-sample states ...?**\n\n**Response:** Before answering this concern, please note that modeling OOD samples is not our ultimate goal, and the reasons are as follows: 1) encountering out-of-sample situations is almost inevitable for any agent due to the presence of approximal errors and other factors (e.g., finite sample size),  2) there exists experimental evidence showing that even when the OOD states are not well modeled (e.g., using our naive Gaussian noise injection), state of art performance can still be achieved if we can find a way to properly guide the agent to navigate to more safe regions (e.g., using our inverse model). Please see our response to Reviewer tZMd for more details. \n\nDue to the above reasons, the focus of our work is actually not to explore the real distribution of out-of-sample states, but on how to improve the robustness of the agent against unfamiliar states, and to encourage the agent to learn more smart behavior when encountering such situations. \n\nAs noise injection is a well-known method to enhance the robustness of the learning algorithm, per your suggestion,  we have explored alternative way to better model the OOD states, which takes into account the physical and environmental constraints, as well as the currently learned policy and Q-values. Particularly, we employ a series of adversarial attacks, including random, action difference, and minimum Q attacks, to generate OOD samples, and use them as data augmentation. The detailed results are given in our response to **W1** of Reviewer tZMd.\n\nFrom the results, we can see that the utilization of a more intricate and meticulously designed noisy dataset is indeed useful in enhancing the performance of our method, probably due to the fact that such methods can improve the efficiency of OOD sampling and expand the coverage of noisy dataset.\n\n**Q2: ... ensure that the constructed transition data accurately represents the true transition ...**\n\n**Response:** Thank you for your comment. Instead of treating it as a dynamic model in representing the true transition distribution, in our work, we regard the inverse model (IDM) as an extended policy. It takes the current state (which may be unseen) $s$ and the one-step target state $s'$ as inputs, providing guidance to the agent on how to reach $s'$ from $s$ as effectively as possible. Although there is no strict theoretical guarantee that the IDM accurately predicts the true transition distribution at unseen states in real-world settings, our experimental results, shown in Fig. 3, 10, 11, and 12 of our paper, demonstrated that the IDM, acting as an extended policy, does aid the agent to behave better under the unseen states.\n\nTo further validate the opinion above, we conduct a simple experiment using MuJoCo suites. State pairs $(s, s')$ are sampled from the offline dataset, and out-of-sample states $\\hat{s}$ are generated via adversarial attacks [1]. By evaluating the average distance between the target state $s'$ and the actual state $\\hat{s}'$ obtained by taking actions from the IDM's policy $IDM(a|\\hat{s}, s')$, we assess the IDM's state recovery ability. A comparison with a CQL policy is also performed, yielding the following results:\n\n| |Halfcheetah|Hopper|Walker2d|\n|-|-|-|-|\n|Distance($s$,$\\hat{s}$)|12.23|3.41|7.75|\n|Distance($s'$, $\\hat{s'}_{IDM}$)|**11.79**|**2.80**|**5.49**|\n|Distance($s'$, $\\hat{s'}_{CQL}$)|17.51|3.52|9.08|\n\nFrom these results we can see that compared to traditional policies like CQL, the action guidance provided by IDM effectively reduce the difference between $s$ and $s'$. It is important to note that while IDM does not guarantee accurate prediction of the true transition distribution at unseen states, the information it provides still proves valuable in mitigating such risks.\n\n**Q3: ... concerns regarding the mismatch between the constructed noisy dataset and the actual out-of-sample states ...**\n\n**Response:** To assess the generalizability of our method, we need to first construct a dataset with more realistic out-of-sample states to test the learnt model. For this we employ a modified generative adversarial network (GAN), which is optimized as follows,\n\n$\\min\\limits_G \\max\\limits_D [E_{s\\sim P(s)}[\\log D(s)]+ E_{P_{G}(s')}[\\log (1-D(s'))]] + \\alpha\\cdot E_{P_{G(s)}}H [\\pi_\\beta(\\cdot|s)]$\n\nwhere $G$ is generator, $D$ is the discriminator, $P(s)$ is the real-word state distribution (the dataset) and $P_G(s)$ is the state distribution generated via $G$. $H$ is the entropy function. In words, the above objective aims to generate real-world samples (i..e., consistent with the in-distribution data) that confuse the behavior policy $\\pi_\\beta$ most; hence its output could be considered as kind of realistic out-of-sample states. \n\nIn evaluation, we initialize the MuJoCo environments with the generated out-of-sample states and assess the generalizability of the learnt agent on them. The results, based on the average of normalized scores and recovery rate, are as follows:\n\n| |Halfcheetah|Hopper|Walker2d|\n|-|-|-|-|\n|CQL|20.4/33.8%|36.5/39.1%|13.1/12.6%|\n|OSR|**40.1/69.9%**|**72.0/88.8%**|**43.3/32.4%**|\n\nThe above results indicate that our OSR effectively guides the agent to recover from most real-world out-of-sample situations (nearly 70%) in the Halfcheetah and Hopper tasks, despite the mismatch between the constructed (Gaussian) noisy dataset and the actual (GAN-based) out-of-sample states. Although the Walker2D task seems challenging to it, our OSR method performs significantly better than CQL, where a risk-guiding mechanism like ours is lacking.\n\nTo gain further insights, we provide visualizations of typical out-of-distribution (OOD) states in the three tasks, along with the corresponding trajectories of OSR. Please refer to the uploaded PDF for the visualizations. \n"}, {"review_id": "tNIggGTPTJ", "reviewer": "Reviewer_kCan", "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.", "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? \n\nIn 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).", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "I think the idea and effort to deal with state deviation in offline RL is meaningful and promising.\n\nThe paper is easy to read and understand.\n\nThe 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.\n\nThe paper lacks sufficient ablation study to support the effectiveness of the OSR/OSR-v term.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thoughtful comments. We provide clarification to your questions as below. We appreciate it if you have any further questions or comments.\n\n**W1: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.**\n\n**Response:** Thank you for the comment. We will make the derivation of Eq.11 from Eq.7 more clear in the revised manuscript. In particular,  a sample from the mixed dataset $\\mathcal D_{tot}$ is denoted as $(\\tilde{s}, a, r, s')$, where  $\\tilde{s}$ can be either the original $s$ or a perturbed one $\\hat{s}$ (according to Eq.5 ). \n\nTo implement Eq.7, we remove the expectation w.r.t. s' inside the KL using Monte Calro approximation with the sample number N as 1, i.e., \n\n$\\min_\\pi\\mathbb E_{\\tilde{s}\\sim\\mathcal D_{tot}}\\Big[KL\\bigg(\\mathbb E_{s'\\sim P(s'|s,\\pi_\\beta)}I^{\\pi_\\beta}\n(a|\\tilde{s},s')\\bigg\\|\\pi(a|\\tilde{s}) \\bigg) \\Big|s, \\pi_\\beta\\Big]$\n\n$\\approx\\min_\\pi\\mathbb E_{\\tilde{s}\\sim\\mathcal D_{tot}}KL\\Big(\\frac{1}{N}\\sum^N_iI^{\\pi_\\beta}(a|\\tilde{s},s'_i)\\Big\\|\\pi(a|\\tilde{s})\\Big)$\n\n$\\approx\\min_\\pi\\mathbb E_{\\tilde{s}\\sim\\mathcal D_{tot}}KL\\Big(I^{\\pi_\\beta}(a|\\tilde{s},s')\\Big\\|\\pi(a|\\tilde{s})\\Big)$\n\nTherefore, more strictly, Eq.11 should be written as:\n\n$L_{sr} = E_{\\tilde{s}\\sim\\mathcal D_{tot}}KL\\Big( I^{\\pi_\\beta}(a|\\tilde{s},s')\\Big\\|\\pi(a|\\tilde{s})\\Big)$\n\nwhere $\\tilde{s}, s'$ are in the same couple sampled from the mixed dataset $\\mathcal D_{tot}$. It should be emphasized that we do not calculate the expectation of $s'$ outside of KL.\n\n**Q1: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?**\n\n**Response:** This is an interesting question, and per your suggestion, we have conducted a more comprehensive ablation study on three environments, shown in the Table below, where each row gives the normalized scores of a specific compared ablated method. \n\n| |Halfcheetah-m.-e.|Hopper-m.-e.|Walker2d-m.-e.|\n|-|-|-|-|\n|QL|9.8|0.3|0.2|\n|QL+BC|41.2|44.7|73.6|\n|QL+IDM|47.5|53.9|80.2|\n|CQL|62.4|98.7|111.0|\n|CQL+BC|85.7|111.8|104.3|\n|CQL+IDM(OSR)|**94.7**|**114.3**|**113.1**|\n\nBased on the above results, we have the following observations:\n\n1) the IDM model is useful - we can see that the performance ranking is: OSR > CQL+BC > CQL , where BC denotes traditional behavior cloning (using behavior policy). The IDM can be thought of as a behavior cloning but taking the consequence of an action into consideration when cloning that action.\n\n2) regulating the Q function search space is important for the generalization capability of an offline RL agent - note that if we replace CQL with standard QL in OSR, we have the following performance ranking: OSR>QL+IDM>QL+BC, which shows that the CQL-stlye value function learning is important both for IDM and BC. One possible reason for this is that the agent is trained under the actor-critic framework where both value function and policy play a role, and most importantly, in the setting of offline RL, conservative learning like CQL is critical for suppressing extrapolation error and overestimation, as pointed out by the reviewer.\n\n3) IDM and CQL are complementary to each other - the IDM provides a way to guide the agent to navigate to safer regions, which allows the agent to learn more smart behavior when encountering unfamiliar states - in the sense that the desired behaviors of an agent should be rational (i.e., being less likely to be punished by the objective function of CQL ) not only under normal in-distribution states but under difficult unseen situations as well (this latter point is less studied in current literature) . The IDM can also be thought of as a mechanism to control the training procedure of CQL, such that it will behave less  'overly-conservatively' during learning, potentially improving its generalization capability.\n\n**Q2:In Fig. 6, it seems that the weight $\\lambda$ for the OSR regularization has little effect on the performance?**\n\n**Response:** Thank you for the comment. The reason why this happens may be that the range of the $\\lambda$s we choose is too small. To further evaluate how the hyperparameters $\\lambda$ and $\\beta$ affect the performance, we have attached more results of sensitive analysis based on the normalized score metrics as follows,\n\nHalfcheetah:\n|$\\lambda$\\\\\\\\ $\\beta$|1e-5|1e-4|1e-3|1e-2|1e-1|\n|-|-|-|-|-|-|\n|0(CQL)|62.4|62.4|62.4|62.4|62.4|\n|0.01|58.7|59.3|64.6|62.1|46.5|\n|0.1|54.2|76.4|83.4|63.7|50.3|\n|0.5|75.4|87.9|**94.6**|82.3|33.6|\n|1.0|73.2|89.2|92.4|46.7|34.4|\n|10.0|64.7|67.9|76.5|43.9|32.7|\n\nHopper:\n|$\\lambda$\\\\\\\\ $\\beta$|1e-5|1e-4|1e-3|1e-2|1e-1|\n|-|-|-|-|-|-|\n|0(CQL)|111.0|111.0|111.0|111.0|111.0|\n|0.01|109.3|110.2|111.6|46.3|20.6|\n|0.1|111.4|112.1|111.3|29.1|18.7|\n|0.5|111.5|112.3|112.9|17.1|20.4|\n|1.0|112.1|111.7|**113.0**|17.4|14.5|\n|10.0|98.3|70.8|69.6|22.6|13.3|\n\nWalker2d:\n|$\\lambda$\\\\\\\\ $\\beta$|1e-5|1e-4|1e-3|1e-2|1e-1|\n|-|-|-|-|-|-|\n|0(CQL)|98.7|98.7|98.7|98.7|98.7|\n|0.01|101.1|102.5|104.6|94.1|33.3|\n|0.1|103.2|108.9|112.4|97.6|16.6|\n|0.5|109.4|112.6|**114.1**|83.3|16.8|\n|1.0|108.2|110.1|113.8|84.7|20.1|\n|10.0|101.7|100.8|89.1|72.8|17.1|\n\nFrom the results, we observe that we should be cautious in choosing a $\\beta$ that is not very large; otherwise, it could lead to failure. We remark that it is better to choose the appropriate hyperparameters in the neighborhood of the bold data listed in the table above."}, {"review_id": "0VbiPrGcWC", "reviewer": "Reviewer_cN6c", "summary": "This 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.", "questions": "* 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\n", "limitations": "The 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.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* 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.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thoughtful comments. We provide clarification to your questions as below. We appreciate it if you have any further questions or comments.\n\n**Q1:... more comprehensive ablation study...**\n\n**Response:** Per your suggestion, we have conducted a more comprehensive ablation study, and the results are shown in the Table below, \n\n| |Halfcheetah-m.-e.|Hopper-m.-e.|Walker2d-m.-e.|\n|-|-|-|-|\n|CQL|62.4|98.7|111.0|\n|CQL+BC|85.7|111.8|104.3|\n|CQL+BC(s.)|91.8|111.2|109.2|\n|CQL+BC(s.a.)|88.3|111.4|108.9|\n|OSR($\\beta$=0)|92.3|111.8|110.1|\n|OSR|**94.7**|**114.3**|**113.1**|\n\nFirst, we run our OSR algorithm with the noisy level being zero and compare it with CQL - an OSR baseline without the inverse dynamic model（IDM). The results indicate that our OSR ($\\beta$=0) (the 5th row) outperforms CQL (the 1st row) significantly on the 3 tasks. This demonstrates the significance of the inverse dynamics model (IDM) in helping the agent to navigate to more safe regions. Similar performance improvement can be observed if we replace the IDM model with the behavior policy (see the CQL + BC setting, the second row), but its effect is not so significant as our OSR method. \n\nThen we introduce the state noise into our base model OSR ($\\beta$=0). The results show that this further improves the performance of the OSR by 2-3$\\%$ (last row), and adding noise improves the performance of CQL + BC (the 3rd and the 4th row) as well except on the Hopper task, while the proposed OSR($\\beta$>0) method consistently improves the performance over all the three environments.\n\n**Q2:... normalizing the states ... useful?...**\n\n**Response:** We added the suggested normalization step into our implementation while keeping a consistent $\\beta$ value. The results shown below reveal that overall normalizing the states is beneficial to the performance and is useful in setting the value of $\\beta$ across different state dimensions (the dimension of the tested environments is 17,11,17, respectively). \n\n| |Halfcheetah-m.-e.|Hopper-m.-e.|Walker2d-m.-e.|\n|-|-|-|-|\n|OSR|94.7|**114.3**|113.1|\n|OSR-norm|**95.8**|113.5|**113.8**|\n\n**Q3:... additional subsection ... on prior works....**\n\n**Response:** We will do that in the revised manuscript.  Our method can be used to improve the robustness of the agent against unfamiliar states, which is consistent with the goal of prior works in robust RL and robust offline RL, although in different manners.\n\n**Q4:... compare prior robust offline RL work....**\n\n**Response:** Per your suggestion, we have compared our method with RORL on the out-of-sample MuJoCo benchmark, and the results are given below; for comparison, we also give the results obtained without using the out-of-samples on these benchmarks. From the table, we observe that RORL generalizes well at out-of-sample states with slight noise, but its performance may suffer a lot under large perturbation. For example, while its normalized score on Halfcheetah-OOS-s. is 103.8, its performance drops significantly to 55.6 on the more challenging Halfcheetah-OOS-l. benchmark. On the contrary, our method performs much more stable across all the environments. In the attached PDF file, we give some visualization of the trajectories generated by both methods in some OOS benchmarks.\n\n|| Halfcheetah-OOS-s.|Halfcheetah-OOS-m.|Halfcheetah-OOS-l.|Hopper-OOS-s.|Hopper-OOS-m.|Hopper-OOS-l.|Walker2d-OOD-s.|Walker2d-OOS-m.|Walker2d-OOS-l.|\n|-|-|-|-|-|-|-|-|-|-|\n|RORL-OOS|103.8|79.8|55.6|111.5|89.5|66.5|117.8|82.2|46.5|\n|RORL-w/o OOS|107.8|107.8|107.8|121.2|121.2|121.2|112.7|112.7|112.7|\n|OSR-OOS|94.1|92.7|91.7|113.3|113.2|110.1|111.4|109.2|106.1|\n|OSR-w/o OOS|94.7|94.7|94.7|114.3|114.3|114.3|113.1|113.1|113.1|\n\n**Q5:... comparisons with prior offline RL works that utilize a reverse model....**\n\n**Response:** Firstly, we would like to emphasize the conceptual difference between the inverse dynamics model (IDM) and the reverse dynamics model (RDM) --  IDM $I(a|s,s')$ behaves more like a policy which gives the action distribution while RDM  $R(s|s',a)$ is a generative model which predicts the reverse state distribution. Hence if the dimensions of state space is much higher than that of the action space, it would not be easy to train a RDM model or generate/image a new state, and vice versa for the IDM model. In addition, the RDM model performs a counterfactual query on the possible cause of action,  which can be problematic with a large action space - by contrast, the IDM model directly infers the most likely action that leads to a given consequence.\n\nPer your suggestion, we have compared our method with ROMI [1]  - an offline RL method based on RDM,  and the results are given below. This shows that although, on average, both methods perform comparably across the environments (average score: 68.9 (OSR) vs. 68.2 (ROMI)), our method performs much better in most 'medium' and 'medium-expert' benchmarks (e.g., Halfcheetah-m.-e., Hopper-m., Hopper-m.e., Walker2d-m.e.) than ROMI, indicating the advantage of using IDM if the underlying dataset covering a portion of high-value areas. However, our method may suffer from very noisy datasets (e.g., Hopper-r.), in which case RDM works better. This suggests that it could be beneficial to combine the advantages of both models for an even more robust offline RL.  \n\n| |Halfcheetah-m.|Halfcheetah-m.-r.|Halfcheetah-m.-e.|Halfcheetah-r.|Hopper-m.|Hopper-m.-r.|Hopper-m.-e.|Hopper-r.|Walker2d-m.|Walker2d-m.-r.|Walker2d-m.-e.|Walker2d-r.|\n|-|-|-|-|-|-|-|-|-|-|-|-|-|\n|ROMI|**49.1**|**47.0**|86.8|24.5|72.3|**98.1**|111.4|**30.2**|84.3|**109.7**|109.7|7.5|\n|OSR(ours)|48.8|46.8|**94.7**|**35.2**|**83.1**|96.7|**113.1**|10.3|**85.7**|87.9|**114.3**|**13.5**|\n\n[1] 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."}]

Author rebuttal:
Thank you for all reviewers' thoughtful and constructive comments on our works discussing a significant but overlooked issue, state distribution shifting,  on offline reinforcement learning.

In summary, our response includes the following aspects:

1.**[Efficiency of noise injection and OOD sampling]** Providing additional experimental results with the adversarial methods onto our implementation, referred as OSR-a, and additiional ablation study including the removal of noise ($\beta=0$) and so on. The detailed results and analysis could be found in the response to Reviewer tZMd's **W1** and **W2** and Reviewer tm1v's **Q1**.

2.**[Validity of the use of IDM]** Providing additional experimental results to explore the behavior of the inverse dynamics model (IDM) we use and its guided policy, enhancing the validity of our work. The detailed results and analysis could be found in the response to Reviewer tm1v's **Q2, Q3** and the visualization in the attached PDF file.

3.**[Additional ablation study]** Providing a more comprehensive ablation study (sensitive analysis), clarifying the role of each component in our work. This could be found in the response to Reviewer kCan's **Q1,Q2** and Reviewer cN6c's **Q1**.

4.**[Additional comparison study]** Providing a further comparison study with related works, including ROMI, a representative offline RL algorithm with the reversed model, and RORL, a SOTA robust offline RL algorithm. The detailed results and analysis could be found in the response to Reviewer tZMd's **Q1**, Reviewer cN6c's **Q5**.

5.**[Performing a prior robust offline RL algormthm in OOS MuJoCo benchmarks]** Providing experimental results of RORL, a SOTA robust offline RL algorithm, in the proposed OOS MuJoCo benchmarks. The detailed results and analysis could be found in the response to Reviewer cN6c's **Q4** and the visualization in the attached PDF file.

6.**[Gap between theory and implementation]** Clarifying this concern in the response to Reviewer kCan's **W1**.

7.**[State normalization]** Providing experimental results of our implementation with the state normalization trick. The detailed results and analysis could be found in the response to Reviewer cN6c's **Q2**.

8.**[Format suggestions]** Including Reviewer tZMd's **W3** and Reviewer cN6c's **Q3**. We will do that in the revised manuscript.

We hope our response could address the reviewers' concerns. If you have any further questions, please post them. We are pleased to have further discussions.

Meta-review: This paper addresses the distributional shift problem in offline RL.  The method is based on promoting the pursuit of state recovery principle when selecting actions, namely biasing towards actions that can recover the state distribution of the behavior policy.  Theoretical results are shown to well align the transited state distribution with the behavior dataset when a state is out of sample.  Experiments on MuJoCo and AntMaze demonstrate the superior performance of the proposed method.

This paper is well written, proposing a solid solution to an important problem.  All reviewers, and myself, agree that it is a good addition to the conference. We hope that the reviews are helpful for improving the paper.
Decision: Accept (poster)