SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\pi(a|\tilde{s}) = \mathcal{N}(\mu_{\pi}, \sigma_{\pi}; \tilde{s})$  and the IDM as  $I^{\pi_{\beta}}(a|\tilde{s}, s') = \mathcal{N}(\mu_{I}, \sigma_{I}; \tilde{s}, s')$ , where  $\mathcal{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  $\pi (a|\tilde{s}) = \mathcal{N}(\mu_{\pi},\sigma_{\pi};\tilde{s})$  and the Gaussian IDM  $I^{\pi_{\beta}}(a|\tilde{s},s^{\prime}) = \mathcal{N}(\mu_I,\sigma_I;\tilde{s},s^{\prime})$  so we can transfer (11) as the following SR loss:
$$
L _ {s r} = \mathbb {E} _ {\tilde {s}, s ^ {\prime} \sim \mathcal {D} _ {t o t}} \left[ \sum_ {d = 1} ^ {D} \left(\log \frac {\sigma_ {I} ^ {d}}{\sigma_ {\pi} ^ {d}} + \frac {\left(\sigma_ {I} ^ {d}\right) ^ {2} + \left(\mu_ {I} ^ {d} - \mu_ {\pi} ^ {d}\right) ^ {2}}{2 \left(\sigma_ {\pi} ^ {d}\right) ^ {2}}\right) \right] \tag {12}
$$
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,
$$
L _ {\pi} = \mathbb {E} _ {\tilde {s} \sim \mathcal {D} _ {t o t}, \hat {a} \sim \pi (\cdot | \tilde {s})} \Big [ Q (\tilde {s}, \hat {a}) \Big ] + \lambda L _ {s r}
$$
where  $\lambda$  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^{\pi_{\beta}}(\cdot |\tilde{s},s')$  to our new policy to enable the agent always transits to those in-sample states  $s^\prime$  from no matter in-sample  $s$  or out-of-sample  $\hat{s}$ , i.e.,  $\tilde{s}$  from the mixed dataset  $\mathcal{D}_{tot}$  mentioned in section 4.1. Then the critic loss function for the Q-networks is,
$$
L _ {Q} = E _ {s, a, r, s ^ {\prime} \sim \mathcal {D}} \Big [ \alpha \cdot \left(\mathbb {E} _ {\hat {a} \sim \pi (\hat {a} | s)} Q (s, \hat {a}) - Q (s, a)\right) + \left(r + \gamma \mathbb {E} _ {a ^ {\prime} \sim \pi (\cdot | s ^ {\prime})} Q (s ^ {\prime}, a ^ {\prime}) - Q (s, a)\right) ^ {2} \Big ]
$$
where  $\alpha$  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,
$$
\begin{array}{l} L _ {Q} = \alpha \cdot L _ {s r - v} + \mathbb {E} _ {s, a, r, s ^ {\prime} \sim \mathcal {D}} \left[ r + \gamma \mathbb {E} _ {a ^ {\prime} \sim \pi (\cdot | s ^ {\prime})} Q (s ^ {\prime}, a ^ {\prime}) - Q (s, a) \right] \\ w h e r e L _ {s r - v} = \mathbb {E} _ {\tilde {s}, s ^ {\prime} \sim \mathcal {D} _ {t o t}} \left(\mathbb {E} _ {\hat {a} \sim \pi (\cdot | \tilde {s})} Q (\tilde {s}, \hat {a}) - \mathbb {E} _ {\hat {a} _ {t a r} \sim I ^ {\pi_ {\beta}} (\cdot | \tilde {s}, s ^ {\prime})} Q (\tilde {s}, \hat {a} _ {t a r})\right) \\ \end{array}
$$
where  $\alpha$  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  $\mathcal{D}$ ) of the mixed dataset  $\mathcal{D}_{tot}$ , as Proposition 2 shows.
Proposition 2. When training on the original dataset  $\mathcal{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_{\pi}$  we use in OSR-v is as follows,
$$
L _ {\pi} = \mathbb {E} _ {\tilde {s} \sim \mathcal {D} _ {t o t}, \hat {a} \sim \pi (\cdot | \tilde {s})} \left[ Q (\tilde {s}, \hat {a}) \right]
$$
To sum up, both OSR and OSR-v optimize actor loss  $L_{\pi}$  to update the policy network  $\pi$  and critic loss  $L_{Q}$  to update our Q-networks, and output the learnt policy network  $\pi$ . 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.
Table 1: Results of OSR(ours), OSR-v(ours), SDC, RAMBO, MOPO, IQL and CQL on offline MuJoCo and AntMaze control tasks averaged over 4 seeds. * indicates the average without 'expert' datasets. The top-2 highest scores in each benchmark are bolded.
Task name	CQL	IQL	MOPO	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.5± 3.4
Halfcheetah-m.	44.4	47.4	40.2	77.6	47.1	48.8± 0.2	49.1± 0.4
Walker2d-m.	79.2	78.3	26.5	86.9	81.1	85.7± 0.6	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.-r.	46.2	44.2	54.0	68.9	47.3	46.8± 0.4	48.5± 0.6
Walker2d-m.-r.	26.7	73.8	92.5	85.0	30.3	87.9± 1.1	87.8± 1.0
Hopper-m.-r.	48.6	94.7	42.7	96.6	48.2	96.7± 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	97.7± 0.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±1.9	92.0±0.8
AntMaze-m.-p.	61.2	71.2	0.0	16.4	71.9	66.0±2.5	71.3±1.3
AntMaze-l.-p.	15.8	39.6	0.0	0.0	37.2	37.9±2.7	38.3±2.5
AntMaze-u.-d.	84.0	62.2	0.0	0.0	57.3	74.0±2.8	69.0±2.5
AntMaze-m.-d.	53.7	70.0	0.0	23.2	78.7	80.0±2.7	77.0±2.4
AntMaze-l.-d.	14.9	47.5	0.0	2.4	33.2	37.9±1.5	34.0±2.2
AntMaze-v0 Avg.	50.6	63.0	0.0	22.3	61.2	64.3	63.4
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  $\beta$  and SR weighting  $\lambda$ . 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-slight
Halfcheetah-OOSmoderate
Halfcheetah OOS-large
Figure 2: 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.
Hopper-OOSlight
Hopper-OOS moderate
Hopper-OOS-large
Walker2d-OOS-slight
Walker2d-OOSmoderate
Walker2d-OOS-large
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 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.
Table 2: Results of CQL, SDC, OSR/OSR-v in OOSMuJoCo setting on the normalized return and decrease metric averaged over 4 seeds. The highest score and lowest decrease are bolded.
Task name	CQL score	dec.(%)	SDC score	dec.(%)	OSR score	dec.(%)	OSR-v score	dec.(%)
Halfcheetah-OOS-s.	56.2	10.0	100.9	0.4	94.1	0.6	98.4	0.7
Halfcheetah-OOS-m.	52.9	15.3	99.8	1.5	92.7	2.1	98.4	0.7
Halfcheetah-OOS-l.	42.1	32.5	83.3	17.8	91.7	3.2	79.8	19.5
Walker2d-OOS-s.	81.1	47.3	102.7	2.5	113.3	0.9	110.3	2.3
Walker2d-OOS-m.	77.4	21.6	102.2	2.9	113.2	1.0	110.9	1.8
Walker2d-OOS-l.	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-l.	54.4	51.0	89.2	20.9	106.1	6.2	105.6	6.7
Besides, as a toy example, we visualize the Halfcheetah agent within an expert trajectory produced by the behavior policy  $\pi_{\beta}$  on the offline dataset  $\mathcal{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 equipped 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 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.
Figure 3: The state recovery process via IDM and OSR on the 'Halfcheetah-OOS-large' benchmark. The interval between every two images is 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 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.
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 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.
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.
This work is partially supported by National Key R&D program of China (2021ZD0113203), National Science Foundation of China (61976115).
[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, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
[27] Hongchang Zhang, Jianzhu Shao, Yuhang Jiang, Shuncheng He, Guanwen Zhang, and Xi-angyang Ji. State deviation correction for offline reinforcement learning. In Thirty-Sixth AAAI Conference on Artificial Intelligence, AAAI 2022, pages 9022–9030. AAAI Press, 2022.