Reward-Consistent Dynamics Models are Strongly Generalizable for Offline Reinforcement Learning
Authors: Fan-Ming Luo
Tian Xu
Xingchen Cao
Yang Yu
Venue: ICLR 2024 spotlight (2024)
Keywords: model-based offline reinforcement learning
dynamics reward
reward-consistent dynamics model learning
Primary area: reinforcement learning
TL;DR: We introduce MOREC, which enhances model-based offline RL by leveraging a discovered "dynamics reward" to improve the fidelity of the generated trajectories, improving the previous state-of-the-art performance by a significant margin.

Abstract: Learning a precise dynamics model can be crucial for offline reinforcement learning, which, unfortunately, has been found to be quite challenging. Dynamics models that are learned by fitting historical transitions often struggle to generalize to unseen transitions. In this study, we identify a hidden but pivotal factor termed dynamics reward that remains consistent across transitions, offering a pathway to better generalization. Therefore, we propose the idea of reward-consistent dynamics models: any trajectory generated by the dynamics model should maximize the dynamics reward derived from the data. We implement this idea as the MOREC (Model-based Offline reinforcement learning with Reward Consistency) method, which can be seamlessly integrated into previous offline model-based reinforcement learning (MBRL) methods. MOREC learns a generalizable dynamics reward function from offline data, which is subsequently employed as a transition filter in any offline MBRL method: when generating transitions, the dynamics model generates a batch of transitions and selects the one with the highest dynamics reward value. On a synthetic task, we visualize that MOREC has a strong generalization ability and can surprisingly recover some distant unseen transitions. On 21 offline tasks in D4RL and NeoRL benchmarks, MOREC improves the previous state-of-the-art performance by a significant margin, i.e., 4.6\% on D4RL tasks and 25.9\% on NeoRL tasks. Notably, MOREC is the first method that can achieve above 95\% online RL performance in 6 out of 12 D4RL tasks and 3 out of 9 NeoRL tasks. Code is available at https://github.com/polixir/morec.

Published as a conference paper at ICLR 2024
REWARD-CONSISTENT DYNAMICS MODELS ARE STRONGLY GENERALIZABLE FOR OFFLINE REINFORCE- MENT LEARNING
Fan-Ming Luo, Tian Xu, Xingehen Cao Yang Yu" National Key Laboratory for Novel Software Technology, Nanjing University, China School Artificial Intelligence, Nanjing University, China Polixirai
{luo fm, xut, caoxe, yuy) amda. .nju. edu. cn
ABSTRACT
Learning precise dynamics model can be crucial for offinn reinforcement learning, which, unfortunately, has been found to be quite challenging. Dynamies models that are learned by fitting historical transitions often struggle generalize to unseen transitions. In this study, we identify a hidden but pivotal factor termed dynamics reward that remains consistent across transitions, offering pathway to better generalization. Therefore, we propose the idea reward-consisten dynamics models: any trajectory generated by the dynamics model should maximize the dynamics reward derived from the data. We implement this idea as the MOREC (Model-based Offline reinforcement learning with Reward Consistency) method, which can be seamlessly integrated into previous offline model-based reinforcement learning (MBRL) methods. MOREC learns generalizable dynamics reward function from offline data, which is subsequently employed as transition filter in any offline MBRL method: when generating transitions, the dynamics model generates batch of transitions and selects the one with the highest dynamics reward value. On a synthetic task, we visualize that MOREC has a strong generalization ability and can surprisingly recover some distant unseen transitions. On 21 offline tasks in D4RL and NeoRL benchmarks, MOREC improves the previous state-of- of- the-art performance by significant margin, i.e., 4.6% D4RL tasks and 25.9% on NeoRI. tasks. Notably, MOREC is the first method that can achieve above 95% online RI. performance in 6 out of 12 D4RL tasks and 3 out of NeoRL tasks. Code is available at
INTRODUCTION
Model- based reinforcement learning (MBRL) approaches encompass techniques that harness either an established environment model one learned to approximate the environment, thereby addressing RL challenges (Sutton& & Barto 1998 Luo et al. 2023 Moerland et al. 2023). These models are primarily employed to forecast forthcoming states arising from the execution of actions within the environment. By leveraging models, it becomes possible to evaluate action sequences or policies through simulation, cireumventing the need for direct interactions with the actual environment and substantially curtailing sampling expenses. As result, model-based methods empower offline RL scenarios where agents exclusively operate with a dataset sampled from the environment, devoid of direct access to the environment itself, thus enabling the adoption of range of efficient methodologies (Yu et al. 2020
2021 Rigter al. 2022 Sun et al. 2023).
One can conduct numerous trial searches within an ideal model until the optimal policies for task are discovered. However, it is incredibly challenging to learn uniformly accurate dynamics model solely through initial supervised lcarning the offline dataset. Without additional guidance the models, model errors will inevitably occur. These errors become particularly pronounced in state-action pairs that fall outside the distribution of the offline data. Such instances may arise when
Yaan Yu is the corresponding author.

---

Published as a conference paper at ICLR 2024
dealing with long forecasting horizons or when the learning policies deviate significantly from the behavior policy of the offline data. These model errors have the potential to be erroneously exploited by algorithms, consequently yielding policies that exhibit significant underperformance.
Research on mitigating model error focuses on two main branches: enhancing policy learning to avoid
model inaceuracies and improving model learning techniques. MOPO (Yu et al. 2020) introduced penalizing uncertain rewards and shortening model simulation horizons to reduce policy deviations and minimize out-of-distribution states visits. Subsequent research introduced more sophisticated conservative strategies, such as MOReL's (Kidambi et al. 2020) early stopping and RAMBO's 's (Rigter al. 2022) pessimistic models, while MOBILE (Sun et (0023) enhanced uncertainty estimation for broader policy exploration. On the other hand, non-conservative methods like MAPLE (Chen et al. 202T) aimed at improving generalization through context-aware meta-policy learning, though their success is ultimately limited by the learned model's s capabilities.
For better model learning, VirtualTaobao (Shi et al.] 2019) led with adversarial model learning for improved accuracy, a technique further applied in complex industrial contexts (Shang al. 2021). This adversarial approach has been effective in addressing compounding errors (Xu et al. 2022a b: 2023), enabling longer trajectory simulations. Moreover, adversarial counterfactual model learning CCen et al. (223)) proved adept at learning causal transitions challenging for straightforward supervised learning methods.
This paper firstly focuses on learning better models. To enhance the generalization ability of dynamics models, our goal is to identify the invariance underlying the data. We identify hidden but pivotal factor named dynamics reward, which represents the inherent driving forces of the environment dynamics. For instance, a naive dynamics reward is function that assigns true transitions and 0 to false transitions. We can conceptualize the dynamies model as an agent which maximizes the dynamics reward (Xu al., (2022)). We then propose the idea of reward-consistent dynamics models: any trajectory generated by the dynamics model should maximize the dynamics reward derived from the data. To implement this idea, we learn dynamics reward functioon through an inverse reinforcement learning (IRL) method from the offline data. We finally incorporate the dynamics reward into the policy learning stage by modifying the model-rollout procedure. Specifically, we (i) regulate the model's output to generate rollouts associated with higher dynamics rewards; (ii) terminate the rollout upon encountering states with low dynamics rewards. This method can be integrated into most prior offline MBRL methods, resulting in novel one named MOREC (Model- based Offline policy optimization with REward Consistency).
In experiments, first evaluate MOREC synthetic refrigerator temperature-contro task. We visualize that the learned dynamics reward aligns closely with the accuracy of transitions, even in OOD regions, demonstrating its superior generalization ability. Furthermore, we demonstrate that the utilization of such a generalizable dynamics reward facilitates the generation of high-fidelity model substantial enhancement in policy performance. Additionally, we evaluate MOREC on 21 typical tasks from two offline benchmarks, D4RL (Fu et al. 2020) and NeoRL (Qin) ct al. 2022). The empirical results show MOREC outperforms prior state-of-the the-th--rt (SOTA) methods on 18 out of the 21 tasks. Notably. in the more difficult NeoRL benchmark, MOREC achieves a remarkable 25.9% average performance improvement over prior SOTAS and solves 3 tasks for the first time (0 previously). In-depth analysis shows a strong correlation between the learned dynamics rewards and model prediction errors. Guided by such an instructive dynamics reward, MOREC shows a significant reduction of the model prediction errors even when long rollout horizon is adopted.
rollouts, leading
PRELIMINARIES AND RELATED WORK
Reinforcement Learning (RL). In RL, we consider the Markov decision process (MDP) SSttoon & Barto 1998). described by a tuple (S, A, P*, rank p'i, 7, Po), where S is the state space. A is the action space. p : S x A - A(S) is the true dynamics, plnsk S x A + R is the task re- ward function, 7 E (0, 1) is the discount factor, and Po is the initial state distribution. The value function VT
,,
E,, [DPe 7 (st+i, ax4:)) 8t = s] the expected cumulative rewards when 8 and following π. Similarly, we can define the action-value function as Q* (²(ssa) a) = E, Σ i= task ($141, a++s) st = s, a, = a). The objective of RL is to find policy i that maximizes the expected value under the initial state distribution, i.e., max, E E (80))
starting at state St 00
2

---

Published as a conference paper at ICLR 2024
1. Model Learning
2. Model Simulation
3. Model Updating
Modal Simulation with Trarsition Filurimaa
Offine Fara
oction
Policy
Offi ne Dara
Model
Replay Buffer
Dyrotticn Rewsnr
Replyy Buffbr
Figure 1: Framework of MOREC. Key distinctions between MOREC and the prior offline MBRL are emphasized in red. The main modifications arise in the model learning and model simulation stages.
Offline RI. Offline RI (Levine al. 2020) emphasizes learning from offline dataset without additional interaction with the environment. Here the offline dataset D = { {(s, r ,')) consists of transitions from trajectories collected by behavior policy иy. A primary obstacle faced in this domain arises from discrepancies between the offline data and the behavior policy, which leads to extrapolation errors (Kumar al.) (019). To mitigate these errors, model-free offline RL methods have integrated conservatism, either by modulating the policy or the Q-function of online RL algorithms (Fujimoto et al., 2019 Bai et al. 2022).
Offline MBRL. Building upon the foundational concepts of offline RI, offline MBRL introduces an advanced strategy by leveraging a dynamics model constructed from offline data. Offline MBRL methods typically possess two stages: (i) learn model of the environment from the offline data D; (ii) learn a policy from the model and D. Let the model parameterized by 0 be Po (8 a.). will be trained by oolliiell maximization:
max (m) == ((,,a,")) (log(Pass' [))))s
Without loss of generality, we assume that the task reward function plsk is known since can simply incorporate yusk into the dynamics model if unknown. The strategy of offline MBRL allows for the generation of synthetic data, potentially enhancing adaptability to out-of-distribution states (Yu et al. 2020 Kidambi et al., (2200. However, the learned dynamics model also inevitably suffers from errors for the limited experience dataset (Xu al. 2022a). To mitigate the challenges stemming from model inaccuracies, several approaches integrate conservatism through methods like uncertainty estimation to guide the behavior policy closer to the available dataset (Yu al. (0220). The most recent works focus on designing better conservative strategies (Kidambi et al.) 2020; Sun et al. 2023 Rigter al. 2022) to unleash the full potential of the model. However, these methods are still limited by the static capabilities of the learned model itself. In this paper, present novel approach to enhance the fidelities model rollouts, through guidence reward signals. This method can not only seamlessly integrate seamlessly but also augment previous offline MBRL techniques.
Inverse Reinforcement Learning (IRL). IRI. Ng Russell 2000; Ni et al. 2020) aims to recover reward function from expert demonstrations. A main class of algorithms (Abbeel & Ng 2004 Ho & Ermon 2016 Swamy et al.) 202T) in IRL infer a reward function by maximizing the value gaps between the expert policy and the other policies, the expert policy should perform well under the desired reward function. It is worth noting that previous IRL methods recover a reward function for a policy, In this paper, we apply IRL learn reward function for the dynamics model.
METHOD
In this section, we will delve into the details of how we learn the dynamics reward function and how we leverage this reward function to facilitate generating high-fidelity model rollouts. We begin with an overview of our method in Section B.1 followed by an explanation of dynamics reward learning in At last, we present the complete process of MOREC in Section 3.3
Section 3.2
3.1 OVERVIEW OF MOREC
We aim to improve the fidelity of model rollouts along with the policy improvement. To accomplish this, must uncover the underlying invariance across different data instances. We have identified a crucial hidden factor called dynamics reward, which represents the intrinsic driving forces of the environment dynamics. By considering the dynamics system as an agent (Xu et al. 2022a). can
3

---

Published as a conference paper at ICLR 2024
assume that the dynamies system was also learned by maximizing the dynamics reward. Therefore, regardless of the policy used to generate interaction data, the data should exhibit consistent dynamics reward.
Based on this idea, we develop a new offline MBRL framework MOREC, which is illustrated in Fig. As general framework, MOREC can be applied to the most existing model-based offline RL methods such as MOPO (Yu et al. 2020), MOREL (Kidambi et al. 2020). and MOBILE (Sun) al.] 2023). Compared with the framework of prior model-based offline RL methods, MOREC incorporates two significant algorithmic designs, which are marked by red color in Fig. T First, we apply IRL infer a dynamics reward function tailored for the dynamics model. We elaborate on the training process of the dynamics reward in Section 3.2 The second algorithmic design is located in the model simulation stage. In particular, we utilize the learned dynamics reward to help generate
high-fidelity model rollouts. The detailed generation procedure is explained in Section 3.3
3.2 DYNAMICS REWARD LEARNING
In this part, we present how we learn a reward function specially designed for dynamics models. Our method is motivated by the new perspective of dynamics models. That is, can view the dynamics model as an "agent" with the new augmented state space S x AA and the new action space S. Accordingly, the behavior policy H is regarded dynamics reward function pD S A x S - R that can judge the performance of a dynamics model. For instance, a naive dynamics reward is a function that assigns 1 on true transitions and O on false transitions. Under the dynamics reward, the true dynamics is good "policy" that achieves the maximum dynamics reward. From this perspective, we can formulate the problem of dynamics reward learning as an IRL problem, which aims to infer a reward function from expert demonstrations collected by the expert policy. In particular, the true dynamics is viewed as the expert policy and the offline dataset collected in the true dynamics is viewed as the expert demonstrations in IRL.. Our target is to learn a reward function that can explain the true dynamics (i.e., the expert policy).
the "dynamics model"". Besides, there is
Based on the above IRL formulation, we leverage the advancements in IRL to help infer the dynamies reward. In particular, inspired by the well-known IRL. method GAIL (Ho & Ermon 2016), we consider the following dynamics reward learning objective.
e(D, =E(...."'> [log( (ss, a, s')]] + F(s.))D,
[log(1 - D(s, a, s'))] (2)
With slight abuse of notation, use D to denote both the state-action and state-action-nexte next-state distributions from the offline dataset. Here P: S x A + (SS) is an arbitrary dynamics model and D ::S S x A x S + (0, 1) is discriminator, which is used to induce the dynamics reward. More concretely, when we consider the dynamics reward of y 7 a, s') = - log(1 D(s, a, s')) (Ho & the difference between rewards the true dynamics p and on the dynamics model P. Notice that a desired dynamics reward should induce large reward difference E(D, P) since the true dynamics achieves the maximal reward such
Ermon 2 !(D, P) can be roughly interpreted
dynamics reward function.
Following the maximum margin IRL method (Abbeel & Ng. 2004 Ratliff et al. 2006 Ho & Ermon 2016), we further define the margin function F(D) := minp EDD P) and consider maximizing such a margin function.
max f(D) = max min E(s,a,')>> [log( ((s, a, s') s'))] + (...)) [log(1 - D(s, a, s'))]. 1)
1> P
In other words, we aim find a dynamics reward that can maximize the reward difference with an adversarial dynamics model P. To optimize this objective, we apply the pradient-descent-aseee method (Lin et al. 2020) which alternates between updating the dynamics reward and adversarial dynamics model. The inner problem corresponds to special case of RL where the horizon equals one, and thus we can utilize the policy gradient method (Sutton et al., 1999) to update the adversarial dynamics model. The outer problem is exactly the binary classification problem by considering the transitions in D as positive samples and the transitions generated by P as negative samples. The detailed procedure is outlined in Algorithm 2 in Appendix B. Note that GAN (Goodfellow et al. 2014) can also be applied in solving the above problem. Here we choose the policy gradient method because it is also applicable in the general RL set-up, which is considered in most IRL works (Abbeel & Ng 2004 Ratliff et al. 22066] Ziebart et al., 2008 Ho & Ermon 2016). Connections between MOREC and adversarial model learning are discussed in Appendix
4

---

Published as a conference paper at ICLR 2024
Algorithm 1: MOREC-MOPO
Input: Offline data D; initialized policy T and critic Qei uncertainty penalty coefficient [;: Invoke Algorithm 2 to learn the dynamics reward rD:
Learn model ensemble P = {po, (s) [M] via Eq !
Initialize an empty buffer B Ø;
for i 1,...,.... Nue do
for j j= 1,..... Nuvlouu do
Sample state 8, from ID uniformly:
Generate rollout ( {st, az. rf ....J following Eq.44) with the dynamics reward r:;
Minus the rewards with uncertainty penalties r2 <- r( BU((P Se, at 9t+1), Vrr E
Insert {s,, at rt, 11.. } to B;
10
Draw samples from D U B and apply SAC to update Ts and Qu:
We theoretically justify the dynamics reward learning algorithm by providing a convergence analysis. Proposition 1. Consider Algorithm
2 Assume that the gradient norm bounded, i.e., Vt [T), VV pl( D,, P.)ll2 S G I Dr, Pill2 S G If we take the step size of {n! =
(ISI t), th A/ D) then have
2 max f(D) 3GD / D
|| T
V |SSA T
T E D.)
3G
Please refer Appendix B.2 for thorough proof. Proposition| | iindicates that the proposed dynamics reward learning algorithm can converge to the global optimum with rate of ((1vVT) Besides, [ laddresses the tabular setting with tabular discriminator and dynamics models. This can be generalized to the linear function approximation setting (Ayoub al. 2020 Cai et al.) 2020). as detailed in Appendix B.3 Finally, we remark that Proposition I shows an average-iterate convergence (Nesterov) 2003). To consist with the theory, practice, we maintain an ensemble of discriminators in the training procedure and output the dynamics reward of
Proposition
T
Σ Dilera.s) ((s,a,5') Sx
s') log 1
3.3 MODEL-BASED OFFLINE POLICY OPTIMIZATION WITH REWARD CONSISTENY
In this part, we elaborate on how MOREC incorporates the dynamics reward for guiding the dynamics model in generating model rollouts. The initial learning of dynamics models in MOREC is the same as that in the prior offline MBRL framework. In particular, we learn a model ensemble, denoted as {Pe... (s'ls, a), Vm E [M] where fm represents the parameter of the m-th model, M symbolizes the ensemble size and [M] == 21, ... .M.. Each model in this ensemble is learned through supervised learning using Eq.
When using the learned dynamics models to generate samples, MOREC incorporates the dynamics reward select high-fidelity transitions from all sampled candidates, which the key difference from other existing model-based offline RL methods. More specifically, at time step I, MOREC begins by sampling from each transition distribution Po. (., |st, a.) N times to obtain M N candidate of next states. The selections then made based a probability distribution induced by the softmax function, which favors transitions with higher dynamics rewards. The softmax distribution also incorporates uncertainty into the model rollout process, which has been previously found to be helpful for policy generalization (Chua et al. 2018). The detailed rollout generation process is:
at TT(--s)); min
Po.. (1s),,)), Vm E [M], n [N]: St+1
8 sm".""
exp(r (St, at s1.7)/k)
(4)
where m² .n* ~ P(). P(m,n) = Em EWEIN exp(r P at s mm' )/x)
Here To is the current learning policy with parameter ó and k is the temperature coefficient of the softmax function. The above process is repeated for H times to generate trajectory with the
5

---

Published as a conference paper at ICLR 2024
Dynamica Reward the Model Transiions dynemica
Dyramis Reward of the true Transiions
(9) Enwwroomenttooo Dynamics temperature
(b) Model Transtions lempenttre
(c1 Mode Tranetion MAE wepperstum
dynamics
2.8
10.02.8
10.02.8
2.8
2.8
2.0
-.0
2.1
2.1
2.1
25
3.8
-3.0
1.5
1.5
1.4
1.4
0.7
0.7
0.7
0.7
0.7
-2.5
-25
0.0
0.0
0.0
0.0
0.5
0.0
0.5
-2.5
-2.5
3.8
3.8 temperature
3.8 temperature
temperature
temperature
temperature
behavior policy
test policy
offline data
Figure 2: Illustration of the (a) environment dynamics, (b) dynamics model, (c) model transition MAE, (d) dynamics reward of the model, and (e) dynamics reward of the true transitions in the refrigerator temperature-contrr control task. Panel (a) visualizes the relationship between temperature, action, and the next temperature, mapped to the 2; y, and color axes, respectively. A point coordinate (r, y), colored by value c, indicates that applying action y at temperature x results in a next temperature of c. For panels (b)-(e), the meaning of the color axes is clarified within the respective figure titles.
maximal length Hoolous- Besides, in order manage cases where the chosen transitions deviate significantly from the true one, leading to an extremely low dynamics reward, early rollout termination also implemented when rd (st, at 11+1) < r T where 7 is a preset hyperparameter. With the above transition filtering technique, we can leverage the dynamics reward to generate high-fidelity transitions, which benefit the subsequent policy learning phase.
As for the policy learning phase, the developed framework MOREC can incorporate the policy learning techniques the most existing offline MBRL methods. In Algorithm [ we provide an instantiation MOREC-MOPO, which utilizes MOPO (Yu et al. 2020) as the policy update component. The uncertainty estimation in line Sis is elaborated in Appendix D.2 In experiments, we also apply MOREC on the SOTA model-based offline RL method MOBILE (Sun et al. 2023).
EXPERIMENTS
In this section, we conduct series of experiments designed answer the following questions: Q1: What is the appearance of the recovered dynamics reward? (Fig. 5
Q2: Can the recovered dynamics reward enhance the accuracy of rollouts? (Fig.
6
Q3: Can MOREC facilitate learning policies with superior performance? (Table
3
4.1 EXPERIMENTAL SETUP
In our experiments. we consider a synthetic refrigerator temperature-conte task, 12 locomotion tasks as part of the D4RL benchmark (Fu al. 2020). and 9 locomotion tasks included in the NeoRL benchmark (Qin et al. 2022). Comprehensive experimental details are provided in Appendix E
4.2 EVALUATION ON A SYNTHETIC TASK
We first consider a synthetic refrigerator temperature-control task where we can delve deeply into the performance of MOREC through visualizations. In this task, an agent controls the compressor's power in a refrigerator with the primary aim of maintaining preset target temperature. The agent's s observations are the present temperature, with the action corresponding to the normalized power of the compressor. We depict the dynamics of this system in Fig. 2 (a). which maps temperature, action, and the next temperature to the x, y, and color axes, respectively. We then use linear behavior policy to collect the offline dataset (also shown in Fig. (a) from the environment.
We train a dynamics model with the offline data which shown in Fig (b). To compare the dynamics model and the true dynamics more clearly, we showcase the mean absolute error (MAE) between the next temperatures output by the true dynamics and the dynamics model in Fig. 2 (c). A comparative analysis of Fig. 2 (a), (b). and (c) suggests the model can align closely with the true dynamics where the offline data is dense, but deviate in data-scarce areas. Notably, the low-error area in Fig. (c), denoted in red, mainly corresponds the regions visited by the behavior policy.

---

Published as a conference paper at ICLR 2024
Mode Tronstion MAE the Fitered Model
behsvior policy
lest policy
2.8
peotuue
Table 1: Returns of the be- havior policy and the poli- cies learned with or with- out transition filtering in the true environment.
-25
144
-12
0.7
00
2.5
0.8
38
6.9 10.0
100
rolbud
102
temperature
milcul
model rilisdd
vollor
Method
Return
Figure 3: MAE of the tran- sitions filtered by the dy- namics reward.
Figure 4: 200-step rollouts in model or environment with the behavior policy or test policy (TF: transition filtering).
with transition filtering
75
transition filtering Dehavvior policy
640
270
To answer Q1, we further train dynamics reward according to Algorithm 2 We depict the inferred dynamics rewards of the model transitions and true transitions in Fig. (d) and (e), respectively. A noteworthy observation is the strong correlation between the inferred dynamics rewards and the MAE of the next temperatures, shown in Fig. (c) and (d). The dynamics reward aptly allocates higher values to model transitions corresponding to a lower
the dynamics rewards of the true transitions (Fig. 2 (e)) are higher than the rewards of the model (Fig. (d)), in almost all of the region. This shows that the dynamics reward model can serve as a scoring function to evaluate the truthfulness of transitions, even in the OOD region.
To answer Q2, we use the proposed transition filtering technique detailed in Eq. to improve the transitions of the dynamics model. We present the MAE of the filtered model in Fig. 3 Compared with Fig. 2 (c). Fig. shows expanded low-error region. This clearly demonstrates the superior performance of utilizing the dynamics reward for transition filtering. A direct comparison between the figures is shown in Fig. in Appendix F.T Subsequently, we compare the accuracy of model rollouts generated with and without the transition filtering technique. Here we consider two rollout policies: the behavior policy and another different test policy which is illustrated in Fig. (a). Fig, A shows the true environment rollout and model rollouts. First, we focus on the behavior policy case, which is shown in the left sub-figure in Fig. 4 As the dynamics model is trained on the offline dataset collected by the behavior policy, can give accurate next-state predictions on in-distribution state-action pairs visited by the behavior policy. Thus, both model rollouts with and without transition filtering can perfectly replicate the true environment rollout.
However, in model-based offline RL, our primary concern is the accuracy of model rollouts induced by learning policies, which could deviate from the behavior policy and visit out-of-distribution (OOD) state-action pairs. Thus turn to the test policy case shown in the right sub-figure in Fig. As the dynamics model cannot give accurate predictions in a single attempt on OOD state-action pairs visited by the test policy, the model rollout generated without the transition filtering technique deviates from the true one significantly. Nevertheless, when equipped with the transition filtering technique, the generated model rollout still perfectly matches the true rollout. These results clearly demonstrate the effectiveness of utilizing the dynamics reward to generate high- -ideefiillity model rollouts on OOD regions. This further suggests that the learned dynamics reward could exhibit a superior generalization ability than the dynamics model.
To answer Q3, we train a policy that aims to control the temperature to C with the reward function (temppeature,) = =- |tepperature, + J. In particular, we apply PPO (Schulman et al. (017) to learn policies from model rollouts generated with and without the proposed transition filtering technique. The detailed returns of the learned policies are reported in Ta We see that the policy learned with transition filtering outperforms that learned without transition filtering by a wide margin. These observations validate that utilizing the dynamics reward can substantially boost the performance of policy learning. The same conclusion can also be obtained by comparing the rollouts of the learned policies in the true environment in Fig. g in Appendix F.
4.3 EVALUATION ON THE D4RL AND NEORL BENCHMARK
To assess the applicability of MOREC in complex tasks, validate it on 21 robot locomotion control tasks from the D4RL (Fu et al. 2020) and NeoRL (Qin et al. 2022) benchmarks. Notice that the NeoRL benchmark is more challenging than the D4RL. benchmark as the offline datasets in
7

---

Published as a conference paper at ICLR 2024
Table Normalized average returns in 12 D4RL tasks, averaged over seeds. Solved tasks denotes the number of the tasks whose scores 2 95.0. The previous SOTA methods are underlined.
Task Name
CQI.
TD3-CC
MOPO
COMBO
TT
RAMBO
MOBILE
MOREC-MOPO
MOREC-MOBILE
53.2
hoppet rnd
31.7
25.4
31.9
26,6
7:4
0.0
17.3
0.7
22,8
mictsh-med
73.0
77.9
hopper
62.8
57.0
1/2.5
$4.1
$1.3
$4.0
87.7
$0.0
55.8
hictth-mocbop
22,1
X4.1
76.5
76.4
hoppa-mes-ng
108.3
96.5
106.5
95.8
med
00.7
106.3
90.8
55,0
95.4
108.2
112.1
110.9 0.56
hopper
08.0
110.7
SL.
130.0
58.2
1126
113.3 2- 0.2
111.5 1.1.5.1...
walker-medd
130.1
112.9
101.9
66.7
0.8
115.5
58.2
70.3
62.3
67.7
83..7
82.8
Solved taoks
2/12
2/12
6/12
Table 3: Normalized average returns in 9 NeoRL tasks, averaged over seeds. Solved tasks denotes the number of the tasks whose scores 2 95.0. The previous SOTA methods are underlined.
Name
COL.
TD3+BC
EDAC
310PO
MOBILE
MOREC-MOPO
MOREC-MOBILE
HallChestabsL
38.2
50.0
60.1
0.6
0..2
Hopper-L
6.2
20.8
Waleerte-a
98.5
64.7
63.0
40.2
11.6
05.0
$6.81119
HallCheatah-M
40.0
54.6
52.5
54.9
62.3
77.5
84.1 0:0.6
$6.0
Hoppredi
51.3
64.5
70.3
48.9
1.0
51.1
83.5 8: 3.8
110.228000 0.6
Waller20-33 24
48.7
07.8
58.6
37.6
30.9
00.2
76.6 1.7
71.1
Haltt Ht
73.3
77-8
75.3
51.4
15.0
900-3 +: 2.0
0
Hopper
78.6
76.8
39.5
11.5
Waller2L-Ht
75.8
00.6
70.3
18.0
70.8
66.1
28.5
10.4
Sulved
NeoRL are collected from narrow distributions (Qin et al.) 2022).
Overall Performance. The normalized average returns are reported in Tables and 3 We observe that MOREC-MOPO and MOREC-MOBILE outperform prior offline RL methods on 18 and 17 of total 21 tasks, respectively. On both the D4RL and NeoRL benchmarks, MOREC-MOPO and MOREC-MOBILE respectively achieve substantial improvements over MOPO and MOBILE, which clearly demonstrate the benefit from the dynamics reward on policy learning and thus effectively respond to Q3. In the more challenging NeoRL benchmark, MOREC offers a more significant improvement over the existing offline RL methods. In particular, MOREC-MOBILE outperforms the SOTA method MOBILE in terms of the average return by a wide margin of 15.7, representing an approximate improvement of 25.9%. Furthermore, we also report the number of solved tasks which refer to tasks with normalized returns exceeding 95. In NeoRL, MOREC-MOBILE is the first method that can solve three tasks while previous methods could not solve any of them. We present the learning curves and the results on Adroit tasks in Appendix [.6 and F.8 respectively.
land
(a) model transition dynamics rewards
(b) true transition dynamics rewards
dynamics reward alfforences
(c) dynamics reward differences
dynamics reward
1.25
1.5
aynamics reward
1.5
2.00
1.5
0.75
1.00
110
1.00
-1.0
0.50
0.75
0.75
0.5
0.5
0.25
0.50
0.50
00.761
0.25
0.0
0.25
10-1
10-1
0.00
10-1 model transition MAE
10:-
0.0
model transition MAE
model transition MAE
(a) p (St, Cg, Sc+1)
(b) r (8t, Cs. 8141)
( (c) ((t,, as 8141) 1441)-r1 - P (8t, 8t+1)
Figure 5: The joint distributions of the model transition MAE, and the dynamics reward rT) for both model transitions and true transitions in the wa 1 lee--ed-rreee task. The z-axes (model transition MAE) all denote the MAE between the true transition and the model transition: || s111lh
Performance of the Dynamics Reward. To answer Q1, we design experiments to evaluate whether the dynamics reward learned by Algorithm 2 can assign transitions proper rewards that are well-aligned with their accuracy. In particular, we choose the malker-med-zee task and apply a policy learned by behaviour cloning in the learned dynamics model collect model rollouts
8

---

Published as a conference paper at ICLR 2024
(a) hictahhmd
hfclahhmed
hiclahhmed-rag
(d) hhclhh-med-eee axp
transiiion MAE
10""
10
00
10
10"
10"
00"
10"
10
10"
10'
100
400
50
rollout step
rollout step
rollout step
rollout step
with transition filtoring
without transition filtering
Figure 6: The MAE in terms of rollout steps with or without transition filtering technique.
{(se. de St+1)h with 500 trajectories. Then we record the dynamics reward (st, at st+1) and MAF ||8t+1 - st+1ll III on the collected model rollouts, where ai+1 the ground-truth next state. The joint distribution of {[1s441 - + 111l1,, (st, de; s141))] is plotted in Fig. 5a We observe a strong negative correlation between the dynamics reward and the MAE. The detailed Pearson correlation between the dynamics reward and the model transition MAE is 70.761, which also denotes a strong negative correlation between both variables. Besides, we remark that the RMSE on the offline dataset is 0.63 (refer to Table T2 in Appendix F.5). and thus the transitions with extremely large MAE (e.g., 3.0) could be regarded as out of distribution data. Even on such OOD transitions, the learned dynamics reward still can assign proper rewards that are consistent with their accuracy, demonstrating its generalization ability. These results clearly demonstrate that the learned dynamics reward is an excellent indicator of the fidelities of transitions.
Moreover, we calculate dynamics rewards on the true transitions ( ae (e, 8(+1). The joint distributions of {[]]++11 - &1+1l1,, r (8t, ate +[+1)]} and {[[ls++ st+1lla r (st: Gt, 8;+1)
r (st: ae, 11)} are reported in Fig. 5b and Fig. 5e respectively. We see that the learned dy- namics reward can accurately assign high values to most true transitions. Note that the transitions (st, ae, $141) and (st, at, $441) only differ in the next state. Nevertheless, the dynamics reward gives totally different values, which suggests that it accurately pays attention to the next state transition. In summary, we empirically verify that the dynamics reward learned by Algorithm is capable of assigning transitions proper rewards that are well-aligned with their accuracy, and thus identifying transitions from all candidates. More visualizations are presented in Appendix [.9) Effectiveness of the Transition Filtering Technique. To answer Q2, we conduct an ablation study to verify the effectiveness of utilizing the learned dynamics reward to filter transitions, which is detailed in Eq. Here we choose the hf tasks where there is no terminal state, which allows the rollout horizon to reach the preset value. We rollout the policy learned by MOREC-MOPO in the learned dynamics model. In particular, we consider two rollout processes: one with the transition filtering technique and the other one without the technique. We respectively take these two rollout processes to collect 100 trajectories and show the MAE on the collected trajectories in Fig. 6 On the one hand, when equipped with the transition filtering technique. the generated model rollouts always keep a relatively small MAE as the rollout step increases. On the other hand, without the transition filtering technique, the MAE the generated model rollouts blows up as the rollout step increases. This phenomenon clearly demonstrates that utilizing the dynamics reward to filter transitions can effectively reduce compounding errors, and thus provide high-fidelity model rollouts for policy learning. An extended ablation study can be found in Appendix F.7
5 CONCLUSION
In this paper, we propose the MOREC method based on the idea of reward-consistent models. MOREC learns generalizable dynamics reward which is then used as transition filter in most offline MBRL methods. We evaluate MOREC through extensive experiments. MOREC outperforms prior methods in 18 out of 21 tasks from the offline RL benchmarks. We empirically validate that the recovered dynamics reward is well-aligned with the accuracy of model transitions even on OOD regions, demonstrating its generalization ability. Consequently, utilizing such a generalizable dynamics reward to filter transitions can effectively reduce compounding errors. However, we also noted the discriminator ensemble could consume a significant memory, especially in the tasks with large state and action dimensions, potentially limiting the scalability of MOREC. A future direction is adopting the last-iterate convergence learning methods (Lei et al. 2021 Golowich et al. 2020). which keep only the last discriminator while maintaining similar convergence guarantees.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
The authors would like to thank Dr. Zuolin Tu and anonymous reviewers for their constructive advice improve this paper. This work supported by the National Science Foundation of China (61921006).
REFERENCES
Pieter Abbeel and Andrew Y. Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the International Conference Machine Learning, Alberta, Canada, 2004.
Kareem Amin and Satinder Singh. Towards resolving unidentifiability in inverse reinforcement learning. CoRR, abs/1601 .06569, 2016.
Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-base based offline reinforcement learning with diversified Q-ensemble. In Advances in Neural Information Processing Systems, virtual, 2021.
Alex Ayoub, Zeyu Jia, Csaba Szepesvári, Mengdi Wang, and Lin Yang. Model-based reinforcement learning with Learning, virtual, 2020.
regression. In Proceedings the International Conference on Machine
Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhi-Hong Deng, Animesh Garg, Peng Liu, and Zhao- ran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. In Proceedings of the International Conference on Learning Representations, virtual, 2022.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge university press, 2004.
Qi Cai, Zhuoran Yang, Chi Jin, and Zhaoran Wang. Provably efficient exploration in policy optimiza- tion. In Proceedings of the International Conference on Machine Learning, virtual, 2020.
Xiong-Huu Hui Chen, Yang Yu, Qingyang Li, Fan-Ming Luo, Zhiwei (Tony) Qin, Wenjie Shang, and Jieping Ye. Offline model-based adaptable policy learning. In Advances in Neural Information Processing Systems, virtual, 2021.
Xiong-Hui Chen, Yang Yu, Zheng-Man Zhu, Zhihua Yu, Zhenjun Chen, Chenghe Wang, Yinan Wu, Hongqiu Wu, Rong-Jun Qin, Ruijin Ding, and Fangsheng Huang. Adversarial counterfactual environment model learning. In Advances in Neural Information Processing Systems, New Orleans, LA, 2023.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, Montréal, Canada, 2018.
Chelsea Finn, Paul F. Christiano, Pieter Abbeel, and Sergey Levine. A connection between gen- erative adversarial networks, inverse reinforcement learning, and energy-based models. CoRR, 0os//611.03885 2016a.
Chelsea Finn, Sergey Levine, and Pieter Abbeel. Guided cost learning: Deep inverse optimal control via policy optimization. Proceedings of the International Conference on Machine Learning. New York City, NY, 2016b.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse reinforce- ment learning. CORR, abs/1710.1 1248, 2017.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach to offline reinforcement learning. In Advances in Neural Information Processing Systems, virtual, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proceedings of the International Conference on Machine Learning, Long Beach, CA, 2019.
10

---

Published as a conference paper at ICLR 2024
Sinong Geng, Houssam Nassif, Carlos A. Manzanares, A. Max Reppen, and Ronnie Sircar. Deep PQR: Solving inverse reinforcement learning using anchor actions. Proceedings of the International Conference Machine Learning, virtual, 2020.
Seyed Kamyar Seyed Ghasemipour, Richard S. Zemel, and Shixiang Gu. A divergence minimization perspective imitation learning methods. In Proceedings of the Conference on Robot Learning, Osaka, Japan, 2019.
Noah Golowich, Sarath Pattathil, and Constantinos Daskalakis. Tight last-iterate convergence rates for no-regret learning in multi-player games. In Advances in Neural Information Processing Systems, virtual, 2020.
lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems 27, 2014.
Elad Hazan. Introduction to online convex optimization. Foundations and Trends in Optimization, (3-4):157-325, 2016.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in Neural Information Processing Systems, Barcelona, Spain, 2016.
Alexis Jacq, Matthieu Geist, Ana Paiva, and Olivier Pietquin. Learning from learner. In Proceedings of the International Conference on Machine Learning, Long Beach, CA, 2019.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, Vancouver, Canada, 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in Neural Information Processing Systems, virtual, 2021.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. MOReL: Model- based offline reinforcement learning. In Advances in Neural Information Processing Systems, virtual, 2020.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In Proceedings of the International Conference on Learning Representations, San Diego, CA, 2015.
Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson. Diseriminator-actor-rrii Addressing sample inefficiency and reward bias in adversarial imitation learning. In Proceedings of the International Conference on Learning Representations, New Orleans, LA, 2019.
Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution matching. In Proceedings of the International Conference Learning Representations, Addis Ababa, Ethiopia, 2020.
Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In Proceedings of the International Conference on Machine Learning, virtual, 2021.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy Q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Systems, Vancouver, Canada, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, virtual, 2020.
Qi Lei, Sai Ganesh Nagarajan, Ioannis Panageas, and Xiao Wang. Last iterate convergence in no- regret learning: constrained min-max optimization for convex-cone concave landscapes. In Proceedings of the International Conference Artificial Intelligence and Statistics, virtual, 2021.

---

Published as a conference paper at ICLR 2024
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. CoRR, 0oos/2005.01643, 2020.
Tianyi Lin, Chi Jin, and Michael Jordan. On gradient descent ascent for nonconvex-concave minimax problems, In Proceedings of the International Conference on Machine Learning, virtual, 2020.
Fan-Ming Luo, Xingchen Cao, and Yang Yu. Transferable reward learning by dynamics-agnostic discriminator ensemble. CoRR, 0tt/22206000238,, 2022.
Fan-Ming Luo,, Tian Xu, Hang Lai, Xiong- Hui Chen, Weinan Zhang, and Yang Yu. A survey model-based reinforcement learning. SCIENCE CHINA Information Sciences, 2023.
Thomas M. Moerland, Joost Brockens, Aske Plaat, and Catholijn M. Jonker. Model-based re- inforcement learning: A survey. Foundations and Trends in Machine Learning. 16(1):1-118, 2023.
Yurii Nesterov. Introductory lectures convex optimization: A basic course, volume 87. Springer Science & Business Media, 2003.
Andrew Y. Ng and Stuart Russell. Algorithms for inverse reinforcement learning. In Proceedings of the International Conference on Machine Learning, Stanford, CA, 2000.
Tianwei Ni, Harshit S. Sikchi, Yufei Wang, Tejus Gupta, Lisa Lee, and Ben Eysenbach. f-IRL: Inverse reinforcement learning via state marginal matching. In Proceedings of the Conference on Robot Learning, virtual, 2020.
Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, and Sergey Levine. Variational discriminator bottleneck: Improving imitation learning, inverse RL, and GANs by constraining information flow. In Proceedings of the International Conference on Learning Representations, New Orleans, LA, 2019.
Rongjun Qin, Xingyuan Zhang, Songyi Gao, Xiong- Chen, Zewen Li, Weinan Zhang, and Yang Yu. NeoRL: A near real-world benchmark for offline reinforcement learning. In Advances in Neural Information Processing Systems, New Orleans, LA, 2022.
Nathan D Ratliff, Andrew Bagnell, and Martin A Zinkevich. Maximum margin planning. In Proceedings of the International Conference Machine Learning, Pittsburgh, PA, 2006.
Mare Rigter, Bruno Lacerda, and Nick Hawes. RAMBO-RI: Robust adversarial model-based offline reinforcement learning. In Advances in Neural Information Processing Systems, New Orleans, LA, 2022.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR, abs/1707 06347,, 2017.
Wenjie Shang, Qingyang Li, Zhiwei Qin, Yang Yu, Yiping Meng, and Jieping Ye. Partially observable environment estimation with uplift inference for reinforcement learning based recommendation. Machine Learning, 1 110(9):603-2666
2021.
Bong-Chegg Shi, Yang Yu, Qing Da, Shi-Yong Yong Chen, and Anxiang Zeng. Virtual-taobao: Virtualizing real-world online retail environment for reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, Honolulu, HI, 2019.
Mingfei Sun, Anuj Mahajan, Katja Hofmann, and Shimon Whiteson. SoftDICE for imitation learning: Rethinking off- off-policy distribution matching. CoRR, abs/2106.03155, 2021.
Yihao Sun. OftiineRL-Kit An elegant pytorch offline reinforcement learning library. https com h//github.ccm//ylaaosss/t 1ttp://11///oeeeeae
Kit 2023.
Yihao Sun, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. inconsistency for model-based offline reinforcement learning. In Proceedings of the International Conference Machine Learning, Honolulu, HI, 2023.
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: An introduction. MIT Press, 1998.
12

---

Published as a conference paper at ICLR 2024
Richard S. Sutton, David A. McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems, Denver, CO, 1999.
Gokul Swamy, Sanjiban Choudhury, Andrew Bagnell, and Steven Wu. Of moments and matching: A game-theoretic framework for closing the imitation gap. In Proceeding of the International Conference on Machine Learning, virtual, 2021.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, Long Beach, CA, 2017.
Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments for rein- forcement learning. IEEE Transactions Pattern Analysis and Machine Intelligence, 44(10): 6968-6980, 2022a.
Tian Xu, Ziniu Li, Yang Yu, and Zhi-Quan Luo. Understanding adversarial imitation learning in small sample regime: stage-coupled analysis. CoRR, abs/2208.01 2022b.
Tian Xu, Ziniu Li, Yang Yu, and Zhi-Quan Luo. Provably efficient adversarial imitation learning with unknown transitions. In Proceedings of the Conference on Uncertainty in Artificial Intelligence, Pittsburgh, PA, 2023.
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 Advances in Neural Information Processing Systems, virtual, 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. COMBO: Conservative offline model-based policy optimization. In Advances in Neural Informa- fion Processing Systems, virtual, 2021.
Brian D. Ziebart, Andrew L. Maas, J. Andrew Bagnell, and Anind K. Dey. Maximum entropy inverse reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, Chicago, II., 2008
13

---

Published as a conference paper at ICLR 2024
CONTENTS
Introduction
Preliminaries and Related Work
3
Method
3.1
Overview of MOREC
3.2
Dynamies Reward Learning
3.3
Model-Based Offline Policy Optmization with Reward Consisteny
Experiments
Experimental Setup
4.2
Evaluation on a Synthetic Task
4.3
Evaluation on the D4RL and N Benchmark
5
Conclusion
Extended Related Work
16
Additional Results for Dynamics Reward Learning
17
B.1 Dynamics Reward Learning Algorithm
17
B.2 Proof of Proposition]
17
B.3 Extension of Proposition
to Linear Function Approximation
18
Connections between MOREC and Adversarial Model Learning
20
Implementation Details
21
D.I Detailed Implementation of Reward Learning
21
D.2 Detailed Implementation of the Uncertainty Estimation MOREC-MOPO
D.3 Detailed Implementation of MOREC
22
Experiment Details
22
E.
Detailed Experiment Setting
22
Hyper-parametery
22
3
Baselines
23
Computational Infrastructure
25
Additional Experiment Results
25
Additional Experiment Results in the Synthesized Task
25
MOREC-MU and its Validations
26
F3
Memory Cost of the Ensemble Dynamics Model
27
F.4 Time Cost Analysis
28
14

---

Published as a conference paper at ICLR 2024
F55
Performance of the Learned Dynamics Model
28
Additional Performance Results
28
Additional Ablation Studies
29
Results of Additional Offline RL Tasks
29
F.
Additional Dynamics Reward Differences Visualizations
31
15

---

Published as a conference paper at ICLR 2024
A EXTENDED RELATED WORK
Offline RL. Offline RL studies the methodologies that enable the agent to directly learn an effective policy from an offline experience dataset without any interaction with the environment dynamics. The major challenge of offline RL arises from the discrepancy between the offline experience dataset and the behavior policy's visitation, resulting in extrapolation errors (Fujimoto et al. 2019 Kumar et al.;
2019). Model-free offline RL algorithms (Fujimoto et al. 2019 Fujimoto & Gu, 2021] Kostrikov
et al. 2021 Kumar et al. 2019: 2020 Bai et al. 2022 incorporate conservatism into policy or Q-function of online RL algorithms tackle with extrapolation error.
The challenge arises from the discrepancy between the offline experience dataset and the behavior policy's visitation. This mismatch can lead to a poor estimation of the Q-function of the behavior policy, plagued by extrapolation errors (Fujimoto et al. 2019 Kumar et al. 2019). Consequently, online off-policy RL algorithms fail to be applied in offline RL. directly.
To tackle extrapolation eтror, conservatism is introduced to offline RL algorithms as a common paradigm. Model-free offline RL algorithms incorporate conservatism or regularization into online RL algorithms by preventing the behavior policy from acting in out-of-support regions (Fujimoto jet al. 2019; Fujimoto Gu 202T] Kostrikov et al. (021) or by learning a conservative O-function
for out-of-distribution (OOD) visitation of the behavior policy (Kumar ct al. 2022). without learning dynamics models.
2019
2020
Bai ct al.
Offline MBRL. Offline MBRL algorithms YYu et al. 2020 Kidambi et al. 2020 Chen al. 202
Yu et al. 202T) leverage a dynamics model derived from offline data to enhance the efficiency of offline RL. Benefiting from the additional synthetic data generated by the learned dynamics model, model-based algorithms are more likely to have the potential to generalize in the states out of distribution and perform better in new tasks. However, the learned dynamics model also suffers from errors inevitably for the limited experience dataset (Xu et al.] [002aa] (Janner et al.) (009)). Thus, some works (Yu et al. 2020 Sun et al., [023 Kidambi et al. I [2200] Yu et al.) (2021) also incorporate conservatism into offline MBRL. For example, some methods (Yu et al. 2020 Bai et al., 2022]
Sun et al. 2023) apply uncertainty estimation and trust states with low uncertainty, while some methods (Kidambi et al. 2020; Yu et al.] (022]) try to limit the behavior policy to acting surrounding the experience dataset. However, despite these efforts, the performance of these techniques is largely dependent on the accuracy of the dynamics model.
Inverse Reinforcement Learning (IRL.) IRL (Ng & Russell 2000 Ni et al. 2020 Ghasemipour
et al. 2019 Swamy et al. 2021 is a process that tackles MDPs devoid of reward functions. The aim of IRL is to deduce these reward functions from series of expert demonstrations. Apprenticeship learning (Abbeel & Ng 2004) trains reward by maximizing evaluation margin between the expert and the policy. In MaxEnt IRL algorithm, the reward is modeled as a maximum likelihood problem by introducing maximum entropy objective (Ziebart et al., 2008). GCL (Finn et al. 2016b) extends MaxEnt IRL to high-dimensional state-action space. Recently, as a special variant of IRL, adversarial imitation learning (AIL) (Ho & Ermon 2016 Finn et al.) 2016a| Fu al., 2017 Kostrikov et al., 2020
Sun et al. 2021 Kostrikov et al., 2019) learns policy and reward in full ad verearrial manner.
Generalizable reward learning (Fu al. (017)), which aims at learning reward functions robust to the dynamics changes, have attracted lots of interests recently. Generalizable reward learning enable the learned reward be reusable in mynsmics-mismatch environment, thus largely extending the scope of IRL applications. Prior works have shown such generalizable reward can be recovered
under some assumptions (Geng et al.) 2020 Amin& Singh, 2016 Jacq et al.) 2019), e.g. assuming the true reward is only related to the state (Fu et | 2017). Besides, |Peng et al. (2019) also show that, by incorporating regularization like mutual information constraints into reward learning, we
can largely improve the generalizability of the reward functions. In this paper, we also introduce regularization to enhance the generalizability of the dynamics reward.
Despite previous efforts, the performance of offline MBRL algorithms is intrinsically limited by the prediction accuracy of the learned dynamics model itself. In this work, with the aim of obtaining more accurate predictions of dynamics, we employ a better generalizable reward of dynamics model, which is learned by IRL, to correct the transitions during the policy learning phase in an offline MBRL algorithm.
16

---

Published as a conference paper at ICLR 2024
B ADDITIONAL RESULTS FOR DYNAMICS REWARD LEARNING
B.I DYNAMICS REWARD LEARNING ALGORITHM
Here we present the completed algorithm for dynamics reward learning. Let W= = (D E R O < D's.a,s') 1. 1, s') 8 x A S} and p = {P | P(s'ss, a) > (s, a,,), Es p(s"ss,a) a) 1, V(s, a) E x A} denote the sets of all feasible discrimi- nators and dynamics models, respectively. Besides, for set X _ R2, we use IIx to denote the fz-norm-based projection operator onto the set X, i.e., (x = argmin || - -y
Algorithm 2: Dynamics Reward Learning
Input:
Offline data D; number of iterations T; step sizes {nP: 778 /rm Initialize the discriminator D1 and dynamics model P1;
for ( I,...,T do
D441 c Tlw
Dr + ne (B(e,,,)")> VD log (Di(s, 8 +
((saa))D,ds VD log (1 - D2(8, a, 8 a',))
Pit1 + 1Ip (Pr - B.a.) u log (s |s, a)) log (1 s((s)
Output: The reward model: (s) -ogg Σi= log(1 Di(s, a, (T)
B.2 PROOF OF PROPOSITION
To prove Proposition] 1 we need the following auxiliary Lemma.
Lemma 1. Consider the objective E(D, P)
B( [Dog( a,))
P.eaa)) [log(1 ((s,,,s))) ((D,)) is concave in D E W and convex in P E P. where W = {D E RISPIAI 0 < s's,a,s') < a,, E S x A x S and P (PERISPLA) P(s'(s, a' P(s'ss,a) = 1, (() E A}. Besides, VD, D' E
- D'll2 x AOSSSIAA VP, P' EPP, P P'Il2 S 2//
Proof of Lemmal 7 First, apply the second-order condition (Boyd & Vandenberghe 2004) to verify the concavity of E(D, P) with respect to D. In particular, we caleulate the gradient v VD(D, P) e RISPAA with the element
V ) aiss, a.), D(s,a,s') a(s,a) a(s.a))((/(s,,)) - D(s.a.s')
With slight abuse of notations, we use diss, a,) 8') and d({s,a a) to denote the empirical distributions from the offline dataset DD Concretely,
Mss (is,a,s)) s') E(enan Tfst 8, a, 8111 = 5 Asssa) E( I[st s, a, = a)
D
[D]
We further calculate the Hessian matrix VED P) RISPLAIX|S PIAI, We note that V)
a diagonal matrix with the principal diagonal elements of
Vileas ((D.P) -d 5') D(s,a,8'
a) (1- D(s, a, 8'))² < 0.
Therefore, the Hessian matrix V',(D,)) E x AAl is a negative definite matrix. Based on the second-order condition (Boyd & Vandenberghe 2004), have that E(D, P) is concave in D. Besides, for any fixed E W. W,D,D P) is a linear function with respect to P. Therefore, ((D,P) is convex in P.
17

---

Published as a conference paper at ICLR 2024
Furthermore, DD E W, we have that ||D D'Il2 1 2 2|IDI|2 S / S|PA Similarly,
we have that VP, P W, |P ll S 2 ||ll = 2 E(aajEss Es a( S
2 (xa))sxa (E. P(s'|s, o))² 2//SiA We finish the proof.
Now we proceed to prove Proposition 1
We have that
+ ΣΣD! T Dr
1
1 T
1 T
T E(D,, P).
min & PEP
Dr, P
2
min PEP
t=1
The last inequality follows that E(D, P) is concave in from Lemmal Notice that the variable P takes the projected gradient descent updates with respect to the sequence of functions {E(D;, P) 1==1
From Lemma | we have that Vt E [2], P(Dr, P) is convex in P and vP, P E P, P. - P S
2 |SSIAA Furthermore, suppose that the gradient norm is bounded, i.., /P Pi)l2 S G
We can apply Theorem 3.1 in (Hazan 2016) to obtain
T Σ min PEP
1 T
S
((D,, 2
((D Pt) 3G" V -
Then we arrive at
1 T
S||A
Dr
((D., P
T
Similarly, the variable D takes the projected gradient ascent updates with respect to the sequence of functions {E(D, pT)]T= From Lemma] we obtain that Vt E [T], ((D, Pt) is concave in D and
VD, D' D, || D'112 < 2 Besides, we suppose that the gradient norm is bounded, i... [D, P S GD. Applying Theorem 3.1 in HHzan 2016) yields that
Σ
(DD,
T Σ -D,P))3GG 3G
/
SS IIA
max DEW
T
T
We combine the above two inequalities and get that
T
T
|SSAA|
3GD /
/
f
Σ D.
max DEW
Σ (D,P)) - 3G
V
T
T
T
|S|AA]
|STLA
max min ((D,)) - -3G
-3GD
DEW PEP
T
T
S
S|2|
A
max DEW
f(D)
3G
T
T
We complete the proof.
B.3 EXTENSION OF PROPOSITION
D TO LINEAR FUNCTION APPROXIMATION
Proposition I considers the tabular setting where the discriminator and dynamics model are directly represented by tables. In this part, we extend Proposition | to the linear function approximation setting. In this setting, based on the linear mixture MDP formulation (Ayoub al., 2020: Cai et al.. 2020), we consider that the discriminator and dynamics model are represented by linear functions. Formally, we have that
V(s,a,s') s') E S x Ax S, Ds a,
(s"ls,, (6(s,a,8').)) 8'),
Here Φ: : S x A x S - Rd is the known feature function and w, E Rd are parameters to be learned. To ensure that the parameterized discriminators and dynamics models are feasible, we consider the following parameter spaces.
W = {w E ERd: ||l|| < Vd.
0< (D(s,, < <1
18

---

Published as a conference paper at ICLR 2024
u E Rd "UER :11u112 S Vd, (ss,a ESxAx8, 8. 2 0,
a)) }1
H'ES
Following (Ayoub et al. 2020 Cai et al. 2020 we consider that the f--norm of parameters is
bounded by vd. In this linear function approximation setting, we consider the optimization problem in the parameter space.
max min f(w, v) == ) [log( (s, s s()o + wEW NEV
[log(1 - Dw s')
Accordingly, we can define the function f(w) = min Aww
For the algorithm, we still consider the gradient-descent-sseen method but the optimization is performed in the parameter space instead of the function space in the tabular setting.
Algorithm 3: Dynamics Reward Learning with Linear Function Approximation
Offline data D; number of iterations T; step sizes {n/ vy Tt=1 Initialize the discriminator parameter w1 and dynamics model parameter for /. ,,...,7 do
Input:
T
W+41 < 1I We 1/f (E(,,,,'') V log (Du, (s,a,
Di (t1s,a) (+1s,a)
V log (1 Das, (s, a, }]])
V log (Pr, L (s' s, a)) log (1- (s,,,,)))))) Output: The reward model: r(s, 8') log(1 (1 E log(1 - Dwi (s,
Pt+1 f Hy (n
Proposition 2. Consider Algorithm Let + pPxd be the feature matrix that aggregates the feature vectors of transitions in D. Assume that rank(() d. Further assume that the gradient norm is bounded, i.e., Vt E [T], (, ve)ll2 2 I G", [[V vi) S G If we take the step size of {nj (d"t,,,, me = (d/(G then we have
T
d 2 max f(w) 30" V NEW
d
3G"
T
Proof. The proof idea is similar to that in Appendix B.2 First, we prove that the objective E(w, v) is concave-convex with respect to (w, In particular, we can calculate the Hessian matrix of Elw, 1)) regarding w.
Velwwv)- = -B -( s7 a') (s, as)
D. (1 DE a, (
1
For any non-zero vector x E Rdd we have that
Vellw, -E, -Es,,,,/)>D
s' (0 (s,,,,'))))
S [1 D2 D' (o ( (s, (3)))
Notice that we consider the under- parameterization case where rank (+)) d. With the rank theorem, we have that (())()()))== O, where A (( denotes the nullspace of कb. This implies that there exists (s, a, s) E D such that (s,, a, s')x * 0. Therefore, obtain that V2((w, v)x indicating that Vll v) is a negative definite matrix. With the second-order condition of concavity, we obtain
'erre we still use the notations and denote the functions in the parameter space when the context is clear.
19

---

Published as a conference paper at ICLR 2024
that f(w, v) is concave with respect to w. Moreover, it is straightforward to observe that for any fixed w E W, E(w,, v/) is a linear function with respect to v. Thus, E(w, v) is convex regarding v.
Furthermore, it is easy to show that the radiuses of W and Q2 are bounded. That is, Yu,w/ W and Vu, // V. || 2'll2 - 2v || v/I12 I 2v d. In summary, have proved that f(w, 2) is concave-convex regarding (w,, 1/), and the radiuses of W and V are bounded by 2v
With these elementary results, we proceed to prove the final convergence result. In particular, we have that
f
Σ
We
min- NEV
((we.P)
T (w -3G"
V/ d
T ((w, 3G" V T 3G" 1=1
d d T
max WEW T
Here incquality (a) follows the concavity of E(v, w) in v. Moreover, notice that the parameter is updated by performing projected gradient descent regarding the sequence of functions {((wr, v) V Recall that E(wc, v) is convex in and the radius of V is bounded by 2vd. Then we can apply (Hazan 2016 Theorem 3. 1), resulting in inequality (b). Similarly, the parameter takes the projected gradient ascent updates with respect the sequence of functions {E(w, v.), Note that { 1/) is concave
in w and the radius of W is bounded by 2 2/d. We can leverage (Hazan 2016 Theorem 3.1) obtain inequality (c).
Then we get that
T
T ((w, vz) - 3G
30"
d T
Wi
2
max WEW
V T
T
T
1=
max min ((w, v) - -G
d
/ 3ar
///
d
wEW WEV
= =max /ww) - 33 WEW
V T
- 3G
which finishes the proof.
C CONNECTIONS BETWEEN MOREC AND ADVERSARIAL MODEL LEARNING
Adversarial model learning employs the generative adversarial imitation learning framework to learn the dynamics model (Shi al. 2019 Xu al., 2022a). It also uses the objective as (2). The equivalence between adversarial imitation learning and inverse reinforcement learning (Finn et al. 2016a) highlights the shared principle with MOREC, which explicitly learns the reward model. Both methods benefit from small compounding error (Xu et al.) 2022a) and exhibit causality-coniiseent
transitions (Chen al. 2023).
However, there are noticeable differences. MOREC employs an ensemble of discriminators as the reward model, which is more stable and generalizable compared to using a single discriminator as in previous adversarial model learning methods (Luo et al. 2022). Additionally, the filtering mechanism allows supervised and adversarial model learning to complement each other. This is particularly advantageous for large training data, supervised model learning helps MOREC leverage all the available data. In contrast, previous adversarial model learning methods solely rely on generative adversarial learning, which may fail to capture small modes in the data.
20

---

Published as a conference paper at ICLR 2024
D IMPLEMENTATION DETAILS
D.1 DETAILED IMPLEMENTATION OF REWARD LEARNING
We adopt the dynamics reward learning approach grounded on the GAIL framework (Ho & Ermon this process involves two main iterative stages: the discriminator updating stage and the dynamics model updating stage.
2016). As delineated in Section 3.2 and Algorithm 2
During the discriminator updating stage (refer to line in Algorithm 2). the parameters of the discriminator are updated using a single gradient step. This step utilizes both the offline dataset and data procured from the current dynamics model. Conversely, in the dynamics model updating phase, data is sampled from the current model to compute rewards based on log(1 - Dr(s. a, s')). The dynamics model parameters are subsequently refined through advanced policy gradient method, namely PPO (Schulman et al. 2017). To ensure stable learning dynamics between the discriminator and the dynamics model, the updating the dynamics model and the discriminator model is repeated 5 times (d.step g g_step- = 5).
Additionally, beyond the standard two-term discriminator loss as depicted in 22. p we integrate gradient penalty to regularize the discriminator, thus curtailing overfitting on offline data. The final discriminator's optimization objective is represented as:
f(D.)) = ((D,)) (.)_))))
(5)
where 7) serves regularization coefficient.
To further optimize computational efficiency, we cap the stored discriminators to a maximum count of L and archive the discriminator every Il iterations. For tasks in D4RL and NcoRL, the parameters are set = 40 and = 10. However, for the refrigerator task, we use 200 and H 1. Other hyper-parameters used for dynamics reward learning are listed in Table
Table 4: The dyyer-parameen of the dynamics reward learning.
Attribute
Value
Number of training iteration Batch size per PPO epoch
5000
20000
Discriminator learning rate Adversarial dynamics model learning rate
3e-4
Value network learning rate Optimizer
le-3
Adam Kingma & Ba 2015
Hidden layers of the adversarial dynamics model
1256, 256|
Hidden layers the value function
[256, 256]
Hidden layers of the discriminator
[128, 256, 256, 128]
0.75 for HalfCheetah in NeoRL and 0.5 otherwise
D.2 DETAILED IMPLEMENTATION OF THE UNCERTAINTY ESTIMATION IN MOREC-MOPO
The proposed MOREC-MOPO technique commences by training an ensemble of dynamics models, denoted as P = (Pe, (s'|s, a) Po. (s' (s'\s,
[M]]. Each individual model within this ensemble, specifically a), characterized as Gaussian distribution, which is parameterized by neural network g(s, 6.). The output of this neural network is splited into two components: the mean g" (s, a; 6,) and the standard deviation gd (s, a; 0.). Thus, the dynamics model can be explicitly expressed as:
g" Po. (s' ss a) (V-lo" ( (s)).
To estimate the uncertainty, we employ the concept of max aleatoric (Yu et al. 2020). This term depends on the maximum aleatoric error. In mathematical context, the aleatoric uncertainty for each model is defined as the L2-norm of its standard deviation, represented as l|g (s, f,)]12 for the i-th model. Therefore, the maximal aleatoric error is formalized (q.
((P, St, des 8t41) ) = max ||g a; 0,)
(6)
21

---

Published as a conference paper at ICLR 2024
D.3 DETAILED IMPLEMENTATION MOREC
We developed MOREC using the OfflineRL -iit codebase (Sun 2023). From this base, we have made the following primary modifications:
1. Initialize the program by loading the pre-trained dynamics reward function.
2. Refine the step method within the dynamics model class:
(a) Execute sampling for each model within the ensemble N times.
(b) Compute the dynamics rewards for the obtained samples.
(c) Formulate a softmax distribution and sample an index from this distribution.
(d) Return the next state corresponding to the sampled index.
E EXPERIMENT DETAILS
E. DETAILED EXPERIMENT SETTING
The refrigerator temperature-control task. In the refrigerator conperature-contry task, an agent controls the compressor's 's power in a refrigerator with the primary aim of maintaining a set target temperature. The agent's observations is the present temperature, with the action corresponding to the normalized power of the compressor.
The system employs transition function f(tem,, a, 2) R x R x (0, + IR that operates in two distinct modes. The modes indicate the current state of the refrigerator door: open (z = 1) or closed (= = 0). These modes are characterized by different rates of temperature change, with the open-door state causing a more rapid convergence towards room temperature (15 CC). The transition function, as detailed in Eq.(7). formalizes this behavior:
ftemp, aes 21) temp, - (0.02 (.006))(15. temp,).
(7)
Fig. (a) graphically illustrates the dynamics of this system when the refrigerator door is closed 0), mapping temperature, action, and the next temperature the x, y, and color respectively.
We generated offline data using a proportional controller as the behavior policy, simulating 2000 interactions in environment where the refrigerator door periodically opens and closes. This policy, formalized in Eq.(8), (8), is also depicted in Fig. (a):
sample
ac temp,) M(a 0.2temp, + 0.75, 0.1).
(8)
In the policy learning stage (Fig. and Table corresponding reward is set to
9 we designed a temperature-control task. The
rtemperaturre -tteperrat + 1/,
where temperature, is the temperature the t-th step.
The NeoRL tasks. We installed NeoRL. from the official repository y and used its 1000-trajectory offline dataset to accomplish all NeoRL. experiments.
E.2 HYPER-PARAMERT
The hyper-parameters for MOREC-MOPO and MOREC-MOBILE derive from the default parameters specified in MOPC and MOBILE in Offlinefa MenneeR--KE Kit Modifying these defaults, extend the rollout horizon from either 1 or 5 to 100, while diminishing the number rollouts per epoch to alleviate computational strain. This adaptation leads to the consolidated syper-parameed for both MOREC-MOPO and MOREC-MOBILE, as detailed in Table
Observing that the model transition MAE consistently exceeds 1.0 when Tn < 0.6 as indicated in we establish y min at 0.6. Notably, this 1.0 surpasses all validation root mean square errors
Fig. 5a
huteas//diibbb.coo/yyaaoo/m/m inekll.iitt
http:///gitthbb..c/myyaaaaam
22

---

Published as a conference paper at ICLR 2024
presented in Table 12 Additionally, value of 0.6 demonstrates efficacy across majority of tasks when applied to MOREC-MOPO and MOREC-MOBILE.
Our investigation primarily revolved around tuning two critical hyper-parameters,, the temperature coefficient A and the penalty coefficient J..
For k. the search space is limited to {0, 0.1), where 0 indicates the selection of the subsequent state relying on the maximum dynamics reward.
Regarding 3, we began by identifying the optimal value for MOREC-MOPO within the D4RL tasks. We determined the search space for 3 by adjusting the default penalty cocfficient using scaling factors (0.1, 0.5, 1.0, 1.0, 2.0, 3.0, 4.03. In the context of NeoRI., given the absence of predefined 33 in the OfflineRL-Kit repository, we need to search the hyper-parameters in the NeoRL tasks from scratch. Setting the base value of / 2.0, we then make search around the base value. It was observed that both hopper-v33 H and Hopper-v3 MM demand a larger 3, whereas Wa Malker2d-v3-d requires a markedly smaller value. The finalized hyper- pareme-e configurations for MOREC-MOPO can be found in T
Transitioning MOREC-MOBILE, we employed a searching strategy akin to that used MOREC- MOPO. Using the 3 values reported by [Sun et al. (2023) as foundation, we conducted multiple searches around this baseline by adjusting 3 using coefficients such as [0.3, 0.4, 0.5, 0.75, 1.0, 1.5). was discerned that MOREC-MOBILA exhibits slightly higher sensitivity to variations in / compared to MOREC-MOPO, This prompted an additional round of maaee--prrmeete refinement for some tasks, based on initial search outcomes. The definitive hyper-parameter settings for MOREC-MOBILE are detailed in Table
Table 5: The common hyper-parameter in MOREC-MOPO and MOREC-MOBILE
Attribute
Value
Actor learning rate
Ic-4
Critic learning rate
3c-4
Dynamics learning rate
le-3
Model ensemble size
7
Number of selected models (M Eq.(4)
5
Number of resample times (N in Eq (4)
2
The number of critics
2
Hidden layers of the actor network
[256, 256]
Hidden layers the critic network
[256, 256]
Hidden layers the dynamics model
[200, 200, 200, 200]
Discount factor 7
0.99
Target network smoothing coefficient Max Rollout horizon Heatou
5e-3
100
Optimizer of the actor and crities Rollout number per epoch Batch size of optimization Batch number of infering reward
Adam
2000
256
4
Threshold of dynamics reward for rollout termination (riii)
0.6
Total steps of optimization
3c6
Actor optimizer learning schedule
Cosine learning schedule
E.3 BASELINES
Here we introduce the baselines used in our experiments, including model-free offline RL and model-based offline RL.
Model-free offline RL.
CQL (Kumar et al. 2020 adds penalization to Q-values for the samples out of distribution;
TD3+BC (Fujimoto & Gu 2021 incorperates a BC regularizaatio term into the policy optimization
objective;
EDAC (An et al.
2021 proposed to penalize based on the uncertainty degree of the Q-value.
23

---

Published as a conference paper at ICLR 2024
Table 6: Particular prree--prameters of different tasks in MOREC-MOPO
TASK
Temperature coefficient
Penalty coefficient |
HalfCheetah-v3-H
2.0
HalfCheetah-v3-L.
2.0
HalfCheetah-v3-M
2.0
Hopper-v3-H
20.0
Hopper-v3-L
2.0
Hopper-v3-M
20.0
Walker2d-v--H
0.1
Walker2d-v3LL
2.0
Walker2d-v3-M
2.0
miiichectah-mddiim----sse
2.5
0.1
0.5
sulfcheetta-mmediume
0.1
0.5
halfcheettal-randdm-vz
0.1
0,5
hopper-medium--e
15.0
10.0
hopper-medium:
0.1
15.0
hopper-random-vv2
10.0
wallker2d-medium-exxprrtt
1,25
walkeer2d-mccuuu--eplaaz-2
0.25
walker2d-mccium--22
0.1
0.5
walker2d-random--2
0.1
1.0
Table Particular sappee-rrameters of different tasks in MOREC-MOBILE
TASK
Temperature coefficient
Penalty coefficient /
HalfCheetah-v3-It
0.8
HalfCheetah-v3-L
0.4
HalfCheetah-v3-MM
0.1
0.5
Hopper-v3-H
0.1
2.5
Hopper-v3-L
0.4
Hopper-v3-M
2.0
Walker2d-v3-H
0.04
Walker2d-v3-L
0.75
Walker2d-v3-M
0.8
halfcheetah- halfcheetah-mmeiim-- medium
medium-expert--v2
0.1
1.0
0.1
0.2
halfcheetah-mediimn-22
0.1
0.2
halfchectah-rannd randetaahaannnnn22222
0,25
hopper- hopper-medium- -exttut--xxprrv-222
1.5
hopper- mediim--cdlmaayyyeee
0.6
hopper-medium-v2
0.1
1.5
hopper-random-v22
0.1
7.5
vellker2d-mediiu--xxxtttzzz
0.2
wallerrad-mediuu--eelaazzzz
0.01
walker2d-medium-y
0.75
walker2d-random-v2
1.0
Model-hased offline RL.
COMBO Yu et al. 202T) which applies CQL in dyna-style enforces Q-values small on OOD
samples;
24

---

Published as a conference paper at ICLR 2024
Model Transition Original Model Color: MAE the next temperature
(6)
Transition the Filtered Model Color: MAE of the nexd lemperallue
2.8
3.7
2.8
3.7
3.1
3.1
2.1
2.1
-25
action
1.4
0.7
0.7
0.6
0.0
-2.5
10.0
-2.5
0.6
3.8
6.9
10.0
temperature
temperature
Figure 7: MAE visualizations of the dynamics models with and without transition filtering. The noises of sub-figure (b) come from the sampling process in Eq. (4).
2022) trains the dynamics model adversarially to minimize the value function without loss of accuracy the transition prediction;
(igterRigter et al.
MOPO (Yu et al. 2020) learns a pessimistic value function from rewards penalized with the uncertainty of the dynamics model's prediction;
2023) penalizes the rewards with uncertainty quantified by the inconsistency of Bellman estimations under an ensemble of learned dynamics models.
MOBILE (Sun al.
TT (Janner et al.
202T) applies Transformer (Vaswani et al. 2017) to modeling distributions over trajectories and uses beam search for planning.
E.4 COMPUTATIONAL INFRASTRUCTURE
All experiments conducted workstation outfitted with Intel Xeon Gold 5218R CPU, NVIDIA RTX 3090 GPUs, and 250GB of RAM, running Ubuntu 20.04.
F ADDITIONAL EXPERIMENT RESULTS
F.1 ADDITIONAL EXPERIMENT RESULTS THE SYNTHESIZED TASK
environment rollout
10.0
with transition filtering
w/o transition filtering behavior policy
7.5
temperature
5.0
target
2.5
0.0
-2.5
0
50
100
150
200
rollout step
Figure 8: Real environment rollouts of the policies learned with and without transition filtering.
MAE comparison between the dynamics models with and without fransition filtering. In Fig. 7 as a complementary of Fig. 3 in the main body, we plot the transition error of the original model transition and the model transition filtered by using the dynamics reward. It can be observed that the low error region has been largely expanded, particularly into OOD region. This clearly demonstrates the superior performance of utilizing the dynamics reward for transition filtering.
Real environment rollout visualizations. We apply PPO (Schulman et al. (017) to learn policies from model rollouts generated with and without transition filtering. The rollouts of the learned
25

---

Published as a conference paper at ICLR 2024
policies evaluated in the true environment are visualized in Fig. g We observe that the policy learned with transition filtering can effectively control the temperature to - Unfortunately, the policy learned without transition filtering induces a temperature that significantly differs from the target. This observation further validate that the dynamics reward can substantially boost the performance of the learned policy.
Ablation studies on the ensemble technique, To validate the efficacy of the ensemble technique, we substitute the ensemble discriminator in (3) by single discriminator. Particularly, we implement a variant dynamics reward consisting of only the last discriminator. In Table g we compare the average model MAE of the original dynamics model, the dynamics model filtered by the ensemble of discriminators, and the dynamics model filtered by a single discriminator. The average model MAE computed over the full state-action space. It can be observed that the single discriminator performs worse than the ensemble of discriminators. This result implies the ensemble technique indeed contributes to the performance and the generalizability of the dynamics reward.
Table 8: Average model MAE over the full state-action space of the original dynamics model, the dynamies model filtered by the ensemble of discriminators, and the dynamics model filtered by a single discriminator (the last discriminator).
Model
Average MAE
Original Model
0.226
Model with Transition Filtering (ensemble discriminator)
0.127
Model with Transition Filtering (single discriminator)
0.141
F.2 MOREC-MU AND ITS VALIDATIONS
In the refrigerator temperature-contro task, we propose a novel strategy that capitalizes on the dynamics reward by integrating it with RL methods. Specifically, we introduce the variant MOREC- MU (Model Update), where the model's s parameters are iteratively updated through RL to optimize the dynamics reward. This process involves the utilization of PPO (Schulman et al. 2017) to maximize the cumulative dynamics reward. In the MOREC-MU framework, a given policy employs the current model to generate batch of rollouts Each transition within these rollouts is then attributed dynamics reward. Following this, PPO updates the model's 's parameters to optimize the accumulated dynamies rewards. This procedure is reiterated over multiple epochs tailor the model to the policy. Notably, we limit the model ensemble size to a single model in MOREC-MU to ensure on-policy sampling.
To ascertain the efficacy of the dynamics reward in enhancing model accuracy concerning a test policy, we updated the model, initially trained by [. using MOREC-MU, employing 60 PPO updating epochs. The updated model's rollout is depicted in Fig. IO It is evident that the rollouts from the MOREC-MU updated model closely align with the environment, underscoring its capability to refine the model for unobserved test policies. Further insights are offered in Fig. [ where the model's ' MAE visualized alongside the test policy. Contrasting with Fig. (c), the adapted model showcases a significant alteration in the MAE map, with the test policy predominantly situated within the low-MAE zones. This implies that the revised model can offer more precise transitions for the test policy.
Table 9:
Mean absolute temperature errors of the learned policies.
with transition filtering
with model updating
w/o transition filtering
behavior policy
Mean Temperature Error Return
0.39
0.37
3.20
1.35
78.0
-74.0
-640.0
-270.0
Lastly, we seamlessly integrate MOREC-MU into the optimal policy learning procedure. We also consider caaperature-contro task, which aims to stabilize the refrigerator temperature at 1'C. We adopt an iterative updating strategy comprising two main steps: (i) refining the policy within the current model to maximize the task reward, and (ii) adjusting the model using the present policy to maximize the dynamics reward. This iterative process is performed for total of 50 cycles.
26

---

Published as a conference paper at ICLR 2024
Color: MAE the next temperature
2.8
2.5
2.1
1.9
action
1.4
0.6
0.0
-2.5
0.6
3.8
6.9
10.0
temperature
test policy
Figure 9: Visualization of the model MAE the model updated via MOREC-MU. The test policy is surrounded by dark red, corresponding low MAE
test policy
environment rollout
env. rollout
10.0
with transition filtering
model rollout with
with model updating
temperature
model rollout with MU
semperatur
7.5
w/o TF & MU
model rollout w/o TF MU
5.0
behavior policy
target
2.5
0.0
-2.5
0
50
100
150
200
50
100
150
200
rollout step
step
Figure 11: Rollouts of the learned policies in the environment.
200-step rollouts in model or envi- ronment with the test policy.
Figure 10:
Specifically, each iteration involves a two-phase update: a 5-epoch PPO update for the policy, followed by 5-epoch PPO update for the dynamics model. The resulting rollouts, generated by the refined policy in the environment, are illustrated in Fig. [1 evaluation of the mean temperature superior performance compared to MOREC, particularly evident when regulating the temperature during instances like the refrigerator door being open. This is further corroborated by the results presented in Table 9 where MOREC-MU yields reduced temperature errors.
Additionally, we provide a quantitative Evidently, MOREC-MU demonstrates
and return in Table 9
Collectively, these experimental outcomes highlight an alternative yet promising approach lever- aging the dynamics rewards, showcasing the notable improvement of MOREC-MU over MOREC. These findings not only advocate for the potential merger MOREC with model updating during policy learning but also emphasize the adaptability of MOREC-MU, which can perpetually refine model parameters alignment with the prevailing policy, thus producing more precise transitions.
F.3 MEMORY COST OF THE ENSEMBLE DYNAMICS MODEL
In Algorithm 2] we utilize an ensemble approach to formulate the dynamics reward function. Given the necessity to store and manage historical discriminators, might anticipate rising memory overhead with MOREC. Nonetheless, as demonstrated in Table [0] the memory consumption associated with the dynamics reward remains relatively modest. Here, the ensemble size is 40.
27

---

Published as a conference paper at ICLR 2024
Table 10: The memory cost of the dynamics reward D4RL tasks.
Task Category Name
Memory Cost (MB)
ha fchettah hopper walker2d
36.843
36,257
36.843
F.4 TIME COST ANALYSIS
To evaluate the computational impact of MOREC, we compared the runtime of MOREC-MOPO and MOPO with identical hyper-parameters. The mean computational time for each epoch in the hal molfcheetah-med. um-v2 task is summarized in Table II Over a span of 3000 training epochs, MOREC-MOPO exhibits an additional overhead of approximately 1.8 seconds per epoch, translating to increase of roughly 12.4% in computational time. Consequently, training MOREC-MOPO policy over 3000 epochs requires an extra 1.5 hours compared to MOPO. Despite this increase, the notable performance enhancement achieved by MOREC-MOPO justifies the additional computational expenditure.
Table 11: The time cost of MOPO and MOREC-MOPO D4RL halfcheetah tasks.
Time Cost (s/epoch)
Task Name
MОPO
MOREC-MOPO
hal fchhetthh-meiiuneee
14.5
16.3
F.5 PERFORMANCE OF THE LEARNED DYNAMICS MODEL
In accordance with methodologies established in MOPO at. 2020) and MOBILE (Sun et al.) 2023), we also train a suite of dynamic models using supervised learning. This is accomplished by optimizing the log-likelihood, depicted in [q.[I. Importantly, prior to training, we partition our dataset into training and validation subsets. The model exclusively utilizes the training subset for parameter updates. For reference, the validation root mean square errors (RMSEs) are presented in Table
We also compare the dynamics models obtained through adversarial training (Alg. 2 with those obtained through supervised learning, and the root mean square model errors of part of tasks are listed [3] It can be observed that the dynamics models obtained through adversarial training typically have higher RMSE error than those obtained through supervised learning. This does not contradict previous research findings on adversarial model learning (Xu et al. 2022a (Shi et al.) 2019). In experiments, did not fine-tune the hyper-paun parameters during our adversarial training and only set the rollout horizon to 1. Therefore, the performance the model obtained through adversarial training is significantly suppressed. Due to the higher model error, only use the model learned via supervised learning for policy learning (Alg.
F.6 ADDITIONAL PERFORMANCE RESULTS
We consolidate Table and Table into Table 14 It is evident from the data that MOREC-MOBILE achieves an average score 80.0, surpassing the prior SOTA results by a margin of 6.7. Notably, MOREC-MOBILE is the sole method reaching the score no less than 80.0. Additionally, it suc- cessfully addresses 9 out of 21 tasks, marking an 80% enhancement compared to the previous SOTA.
The learning curves of both MOPEC-MOPO and MOREC-MOBILE are depicted in Fig. and Fig. [3] These curves indicate that both methodologies maintain stable learning progression. Furthermore, the small shaded region, representing the standard error, underscores the robustness of MOREC against variability in random seed initialization.
28

---

Published as a conference paper at ICLR 2024
Table 12: The model root mean square error (RMSE) standard error on the validation data, averaged over 5 seeds.
Model Validation Error
HalfCheetah-v344
0.5273 0.0108
HalfCheetah--3-LL
0.4164 1 0.0076
HallCheetah-v3-M
0,4917 +: 0.0095
Hopper-v3-H
0.0615 + 0.0026
Hopper-v3-L
0.0878 0.0017
Hopper-v3-M
0.0664 = 0.0010
Walker2d-v3-H
0.4785 + 0.0069
Walker2d-v3-L.
0.5046 + 0.0052
Walker2d-v3-M
0.4488 0.0066
haffchetaah-mcciuu--spperyy
0.4281 + 0.0185
malichetahhmmmdiuu--eelaaa-2
0.7182 0.0077
medicheetah-mmciiun2222
0.4742 + 0.0071
balfcheetah--andoone
0.3449 + 0.0039
hopper-mediimm-eppere
0.0443 + 0.0008
hopper-mmeduum--eelaa-222
0.0734 11: 0.0024
hopper-medium--22
0.0609 0.0027
hopper-random--2
0.0322 + 0.0041
walker2d-meciim--xxppe
0.3361 + 0.0047
walkeer2d-meddiun--rppaay-
0.6256 + 0.0076
walker2ddmmediun-222
0.5745 0.0070
walker2ddraandom--22
0.5963 0.0071
Table 13: Root mean square model error comparisons between adversarial-leaming models and supervised-leariing models.
Dataset
Adversarial Learning
Supervised Learning
walker2d-medium-cxpprr
0.6321
0.3361
walker2dd-mediuu--relaa---
1.3427
0.6256
walker2d-medium--22
0.6589
0.5745
walker2d-random-v2
2.7218
0.5963
Average
1.3389
0.5331
F.7 ADDITIONAL ABLATION STUDIES
In order to validate the effectiveness of the transition filtering technique, in Fig. 6 we implement a variant MOREC, which is not equipped with the transition filtering technique. We set the syper-parameters of the variant to be the same MOREC. We then train policies with both methods the hffctah tasks of the D4RL benchmark. The final policy performance is shown in Table [5] We find the policy performance of MOREC-MOPO dropped lot without the transition filtering technique, implying the significance of the technique.
Besides, we additionally consider the maximum transition MAE. For each model rollout, obtain all its transition MAEs and choose the maximum one for this rollout. Then, we take the average of these maximum MAEs across all rollouts collected during full training process and obtain the maximum transition MAE. We list the maximum transition MAE for 12 D4RL tasks in Table [6] showing significantly MAE reduction MOREC-MOPO. This result implies MOREC-MOPO can avoid inserting the outlier transitions to its replay buffer and ensure the training stability.
F.8 RESULTS OF ADDITIONAL OFFLINE RL TASKS
In addition to evaluating performance on locomotion tasks within the D4RL and NeoRL benchmarks, study extends to the Adroit domain, which has been utilized in prior research, such EDAC (An et al. (021) and MOBILE (Sun et al. 2023). Specifically, we focus on two tasks: men-cloned-vz and pen-human--ue We employ the MOREC-MOBILE method and present the comparative
29

---

Published as a conference paper at ICLR 2024
Table 14: Performances with maximum improvement ratios in D4RL and NeoRL benchmarks. Averaged over 5 seeds. The previous SOTA algorithm is subscripted.
TASK
MOREC-MOPO
MOREC-MOBILE
PREV-SOTA
Max Improvement Ratio
HaltCheetah--3-44
90.27
98.50
83.00 (3.00LE
15.50
HalfChectah-v3-1
53,50
51.22
54.70 MOBILL)
-1.20
HallCheetah-v3-MM
84.08
85.99
(OBBLL)
8.19
H3ppen-v3-H
72.77
108,48
(7.80 (OBBLE)
20.68
H3--er-v3-L.
25.36
26.75
18.30 (OBILE)
8.45
Hopper-v3-M.
83,50
109.23
70. (TD-BC)))
38.93
Walker2d-v3-H
83.04
79.30
74, 74,90 74,90
8.14
Walker2d-v3-L
65.01
56.81
(OQL))
20.31
Walker2d-v3-M
76.59
71.07
62.20 (MOBILE)
14.39
hallcheeuuh-mmchhu--epeeeee..
112.07
110.94
108. 20 (MOBILE)
3.87
balfcheetub-mmciin--rphhe
76.45
76.40
71.70
4.75
halfchectah- memeuuuuaammmin---2222
82.27
82.12
77. MORI
4,37
halfcheetah
saccheetaa-rrndom-22 random-v2
51.57
53.19
39. (9.30 MUBILL)
13.89
hopper- gopper-medium--xxeert-s
113.25
111.47
112. 60 MUBILE)
0.65
105.15
105.45
103. 08,90 (OBILE)
1.55
hopper-medium-
106.96
107.97
106. 106.60 (MOBILE
1.37
hopper-randone v2
32.10
26.58
31.90 (OBILE)
0.20
walker2d- medium--xxperr--22222
115.77
115.52
115. (MOBILE)
0.57
walker2d- met--dd-mmmiiu--eppa-a---.
95.50
95.83
89.90 (MBBILE)
5.93
walker2d medium---2
89.93
85.76
92.50 EDAC)
-2,57
walker2dlrandome
23,53
22.75
16.60 (0BB
6,93
Average
78.0
80.0
73.3
8.3
Solved tasks (performance
6/21
9/21
5/21
het hellkheeesh---eddd---l----------
halfcheelal--mediunn-2
hrlfcheelsh-enddom-22
90
retum
90
60
60
40
30
30
15
0
limestep
106
limeslep
106
limestep hosppr--nedkum-12
1dB
limestep
166
hogper-medim--epdeeee
hxspper-anndom-12
120
120
return
80
90
90
24
60
60
60
30
30
3
0
3
Timestep
106
timeslep
106
limestep walkordd 2.dd-2d-dddum-222
1cB
limestep
1oB
walkor2a-mcdlun-ooooe -2
walker2d-medin wtt.:/d-mmimm--coo replay--2
walker2a Salker2d-random-2
120
100
24
retur
90
75
75
16
60
50
30
timastop
1e6
1e6
timestop
1e8
MOREC_MOPO
MOREC_MOBILE
Figure 12: The learning curves of MOREC-MOPO and MOREC-MOBILE in D4RL tasks. The curves are shaded with standard error over 5 seeds.
results in Table 17 performance relative to both MOPO and MOBILE across these tasks. However, it is noteworthy that in the pen-human does not perform as well EDAC. This performance diserepancy is likely attributable to the reduced size of the dataset, which may constrain the model performance more significantly than in the pen-cloned-vr task.
The comparative analysis reveals that MOREC-MOBILE achieves superior
task, which is characterized by a smaller offline dataset, MOREC-MOBIL
30

---

Published as a conference paper at ICLR 2024
HaffCheetah-v3-H
HalfCheetah-v3-L
HalfCheetah-v3-M
100
45
75
retum
25
15
25
0
2
1
3
3
timestep Hopper-v3-H H
1e6
timestep Hopper-v3-L.
1e6
timestep Hopper-v3-M
1e6
120
30
90
24
return
60
60
12
30
30
6
0
2
3
3
0
3
limestpp Walker2d-v3-H
1e6
timestep Walker2d-v3-L
1e6
timestep Walker2d-v3-M
1e6
80
80
40
20
20
20
3
timestep
1e6
timestep
1e6
timestep
1e6
WOREC_MOPO
MOREC_MOBILE
Figure 13:
The learning curves of MOREC-MOPO and MOREC-MOBILE in NeoRL tasks. The curves are shaded with one standard error over 5 seeds.
Table 15:
The final normalized policy returns of MOREC-MOPO and a variant of MOREC-MOPO without transition filtering, averaged over seeds.
rnd
med
med-rep
med-exp
MOREC-MOPO
51.6 + 0.5
82.3 + 1.1
76.5 + 1.2
112.1=1 1.8
without transition filtering
38.4
2.1
44.2 4.5
73.3 + 1.5
80.4 = 9.2
F,9 ADDITIONAL DYNAMICS REWARD DIFFERENCES VISCALIZATIONS
We analyze the relationship between dynamics rewards and the model transition MAE, denoted as ||st+1 + II,, across 12 D4RL tasks as shown in Fig. 14 Utilizing policies derived from behavior cloning, we sample trajectories from each learned dynamics model, without transition filtering. The joint distribution of ([|1s41 sl+1l,, |l, p (st, Gt, s1+1) pD (st, at, 8t+1))) is depicted in Fig. 14 with (-) representing the dynamics reward, 8141 the next state, and 8141 the next state predicted by the model. It is noteworthy that trajectories in the random data subset are drawn from a random policy, resulting in a tighter distribution compared to other data sources.
In tasks such as 2 and hopper, the discrepancies in dynamics rewards are predominantly positive, underscoring the potent discriminative prowess of the dynamics reward. This indicates that, for these tasks, the dynamics reward consistently assigns higher values to true transitions over model transitions. The distribution patterns for wa kelkdd are distinct from those of hopper. This distinction stems from the variation in one-step model errors between these tasks, as elaborated in Table 12 The root mean square error (RMSE) for walker2d is nearly an order of magnitude greater than hopper. Consequently, cumulative errors in hopper smaller z-axis range. Specifically, the r-axis span for hopper r is roughly [0.001, 0.5], in contrast to [0.008, 7] for wa ker2d. In wall ke rdd, dynamics rewards begin exhibit significant differences
contained, leading to
31

---

Published as a conference paper at ICLR 2024
Table 16: The average maximum transition MAE standard error, averaged over 5 seed, in 12 D4RL tasks. For each model rollout, we obtain all its transition MAE and choose the maximum one for the model rollout. Then, we take the average of these maximum MAEs across all rollouts and obtain the maximum transition MAE.
MOREC-MOPO
without transition filteing
halfchectah-mmcduum-eprrttt
6.066 + 1.777
249906.304 + 48324.466
malficheraa-mmddium-rppaaaya
3.860 + 0.102
11.653 + 5.191
malfchectah-mediimm-e
3.282 + 0.049
304736. 147 + 71894.240
halfcheetah-randoom--
3.137 + 0.252
212908.399 396 + 33037.00 .017
0.556 0.024
1.553 + 0.300
hopper-meddum-relly-yy
0.598 1 0.005
0.613 1 0.014
hopper-medium-v2
0.529 + 0.004
0.581 + 0.018
hopper-random-v2
0.183 = 0.003
0.185 + 0.004
walker2d-
expertt-v2
3.411 + 0.160
8.453 + 0.660
wallker2d-mmeddim--eeplaa--
5.619 0.130
8.568 + 0.570
walker2d-meeiium-22 walker2d-random--2
3.892 + 0.178
4.532 + 0.360
6.964 + 0.601
12.640 + 1.344
Table 17: Normalized average returns in 2 D4RL Adroit tasks, averaged over 5 seeds. The previous SOTA methods are underlined.
BC
CQL.
EDAC
MOPO
MOBILE
MOREC-MOBILE
pen-cloned-vl 38.3
27.2
68.2
54.6
69.0
74.3 1 14.4
pen-human pon-human-vl
25.8
35.2
52.1
10.7
30.1
37.3 + 18.8
when MAE approaches 0.3, which aligns with the maximum transition MAE observed in hopper. This masks the potential trend of increasing dynamics reward discrepancies with escalating transition MAES.
As for the half Salfchhetah tasks, on the contrary, the dynamics reward function sometimes gives higher reward to the model transition than the true transition, which is not desired. We believe that this phenomenon is due to the limited generalization range of the dynamics reward. In some hal Selfcheetahh tasks, when the transition MAE is too large, it exceeds the generalization range of the dynamics reward. The experiment results in Fig. 14jindicaee that, the ha 1 wcheetheetooooo tasks, the generalization range of the dynamics reward is within the scope where the MAE is less than 1.0. order to discover why MOREC also works in the halfcheett
tasks, we additionally visualize 11, (st,, Gg, sx+1))) in Fig. 15 We find the dynamics reward
the joint distribution of [[]]s of hal halffheett ah shows a similar trend as walker2d: As the model transition MAE increase, the dynamics rewards reduce gradually. Because the behavior policy of hal sellfcheetale ah-random--- a random policy, its distribution is narrow and cannot show the aforementioned trend. As a result, the dynamics reward can still give low dynamics reward to the high-MAE transitions. Consequently, the rollout can be terminated on time before the MAE diverges.
32

---

Published as a conference paper at ICLR 2024
driamis ieeard fftreees
1.00
-15
1.00
1.00
0.75
0.75
155
0.75
0.75 0.75
0.50
0.59
110
0.50
0.50
0.50
0:25
0.5
0:25
0.
0.25
0.25
-0.2
10t2
10-1
19a
10t
10-1
10t1
10a
10
-025
10-3
10-1
10"h
101
0.00-0.25
10::
10"
10
model transitonn MAE
transiton
hospae-mectumee?
romand
0.2
0.b
1:00
0.10
0.6
0.75
0.1
0.05
0.500.025
0.03
2
0.0
0.25 0.000
-5
-0.05
0.2
10a1
15-1
15-1
10
10-1
10--
10-1
10-3
10-1
10-1
10-3
10-1
10"1
suffctewcch-mmeddd
bafthhatat-meciameee
1.0
1.0
-5
0.5
0.5
1.0
0.5
0.4
0.3
10
0.0
0.0
0.5
0.0
-02
0.0
-0.
-0.5
-0.5
0.1
-1.0
-1.0
-02
10
10
10
19
30
10-3
10P
10
10
10
10
10"
30
Figure 14: The joint distribution of the dynamics reward differences (se, 88+1)
D (8t, Ot, 8141) and the model transition MAES
|| 8t+1ll ll in 12 D4RL. tasks.
1.4
1.5
1.4
wwk.erd-madur--eee mollayy
14
waber25--redumm222
wallar26eandsmm222
1.0
-1.0
1.0
-1.0
-10
10
0.8
0.6
0.5
10.
0.5
05
10-2
10
10f
10
10-2
10-4
10?
10-
10-2
10-1
10>
10
10 model
10ml
model isttion MAE
model transition
model transiioon MAE
MAE
expert-v2
hoppen-
hoppenrandom-22
1/2
1.0
1.0
1.0
1.1
1.0
005
1.0
0.5
-0.
005
0.9
0.0
0.0
0.9
0.8
10
10-2
10-1
10²
0.00
0.8
10-1
10
10
0.0
0.8
10-3
10t4 mode tonsition M&E
10-2
0.0
08
10-3
1012
3011
0.0
model tra isiion MAE
model transiionn MAE
mrdel transition NA8
1.4
fofthhetah-medum-eeee
1.5
1.4
battcheeth-mmedun opayy92
1.5
1.4
haffcreetah medimmm
15
14
ho tatchheeah-anddome
1.5
1.2
1.2
1.2
1.2
1.0
1.0
1.0
1.0
0.8
0.8
0.8
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.4
10rl
10A
-0.0
...0
10-4
10F
..00
10-4
0.0
model isliin MAE
model ttanstton MAE
modal tianstion NAE
mrodal transition NAE
Figure 15: The joint distribution of the dynamics reward for model transitions p ats 1(+1) and the model transition MAES || - -stall |I in 12 D4RL tasks.
33


Reviews:
[{"review_id": "MGtt1LsDfM", "reviewer": "Reviewer_HqU2", "summary": "Paper describes a new method for offline model-based RL using a new concept of \"dynamics reward\", and present results showing its effectiveness.", "questions": "- It seems to me the dynamics reward model kind of serves as a density model of the training data? If it is, then why not just learn a density model? If not, then what is the difference?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Quality\n- Paper presents a number of experiments showing the effectiveness of the proposed method, including SOTA results on standard offline RL benchmarks\n\nClarity\n- Paper was written clearly, and method and experiment were easy to understand\n\nSignificance\n- Based on the results, I believe this paper will be a nice contribution to the area of offline model-based RL", "weaknesses": "- A major contribution of this paper is the new concept of the \"dynamics reward\", and the algorithm used to learn the dynamics reward model (GAIL). However, it was not clear to me what this dynamics reward should represent. It would be helpful to describe the optimal closed-form solution (if there was one), or include some additional discussion on an intuitive interpretation of it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "J0qHfneRbt", "reviewer": "Reviewer_Gxsz", "summary": "This paper presents an offline model-based RL method that aims to reduce the prediction error of dynamics model in OOD states. The main idea is to learn a dynamics reward which is assigned a higher value when the input transition is close to ground-truth transition. This reward is utilized when predicting the future states with dynamics models, by weighting the future predicted states with their predicted dynamics rewards. Experiments first provide a supporting experiment and analysis on synthetic task and then evaluate the proposed method in D4RL and NeoRL offline RL benchmarks.", "questions": "- Could the authors answer & address my concerns & questions in Weaknesses?\n- What does ``Accordingly, the behavior policy $\\pi_{\\beta}$ is regarded as the “dynamics model”` mean? It's difficult to parse and seems out-of-context as the behavior policy is not even defined yet before this line.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The idea of learning dynamics reward is new to my knowledge.\n- The paper presents helpful analysis and supporting experiments to understand how the proposed method works.\n- The proposed method can consistently improve the performance of offline RL methods in most considered tasks.", "weaknesses": "- The interpretation of dynamics model as an agent and introducing the concept of reward adds unnecessary complexity to the method and makes it a bit difficult to easily understand the method. Simply formulating the main idea with adversarial generative training and using the score D as a reward could make the paper be more simple and easy to understand.\n- The paper introduces multiple hyperparameters and did quite extensive hyperparameters search (e.g., temperature, penalty, and threshold, ..). Making sure that the baseline is fully tuned with the similar resource given to the proposed method could be important for a fair comparison.\n- Figure 2 is very difficult to parse and how this leads to the conclusion that dynamics rewards is superior to dynamics model is not clear.\n- The additional complexity of learning dynamics reward and implementation complexity might limit the wide adoption of the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ReyLuIVNsv", "reviewer": "Reviewer_zEZ1", "summary": "Model-based RL suffers from compounding error as a result of recursive model rollouts. This is especially pronounced in offline model-based RL, where there is no online data to correct model errors. Prior offline MBRL methods typically address this by discouraging the policy from visiting states where the model is uncertain, via some uncertainty quantification. This paper introduces a novel perspective: they learn a \"dynamics reward\" that evaluates the dynamics consistency of transitions and use it to filter model rollouts to only produce trajectories that are plausible under the dynamics. Concretely, they train the dynamics reward model using inverse max margin IRL. Rollout filtering is performed in a beam-search fashion, where at each rollout step they sample a number of transitions from the model and weight them by their dynamics rewards. Applying this method on top of existing offline MBRL methods results in significant improvements over prior state-of-the-art on D4RL and NeoRL benchmarks. Overall, this paper makes a solid contribution to offline model-based RL by introducing a novel perspective for tackling compounding error and distribution shift.", "questions": "- When training the dynamics reward, the (s, a) pair is always sampled from the dataset, whereas the (s') is either from the dataset or from the adversarial dynamics model. This makes the dynamics reward robust to OOD (s') but not (s, a). How does the model generalize to OOD (s, a)? Can you evaluate the dynamics model explicitly on OOD (s, a) pairs and see if the dynamics reward is still higher with valid transitions? \n- How does this method compare to a simple filtering technique like selecting the transition with the lowest bellman error?\n- When filtering, why construct a softmax distribution instead of directly taking the argmax action?\n- How does the reward dynamics relate to energy-based models? It seems quite similar to an energy function on (s, a, s').\n\nSuggestions:\n- Table 1 could use a more detailed caption.\n- It would be nice to use this for model-based offline policy evaluation [1]. I would imagine the reduction in compounding error to benefit OPE a lot.\n\nReferences:\n1. Michael R. Zhang, Tom Le Paine, Ofir Nachum, Cosmin Paduraru, George Tucker, Ziyu Wang, Mohammad Norouzi. Autoregressive Dynamics Models for Offline Policy Evaluation and Optimization. ICLR 2021.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper introduces a novel and credible remedy for a long-standing problem in model-based RL, namely compounding error due to recursive model rollouts resulting in increasingly out-of-distribution inputs. While prior methods mitigate this issue by discouraging the policy from visiting states where the model is uncertain, this paper proposes to filter model rollouts using an adversarially trained \"critic\" function. This approach can be easily combined with existing offline MBRL methods and lead to significant performance gain.\n- The authors provide extensive evaluations of their method on a variety of domains, ranging from analytical toy tasks to standard benchmarks. Their method demonstrates a universal improvement across these tasks.", "weaknesses": "- There are limited ablation studies in the paper. This makes it illusive why the method actually works so well, and when it works well. Please refer to Questions for more details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ge0DdzWqkB", "reviewer": "Reviewer_F8oz", "summary": "This paper proposes a method that constructs a discriminator (known as dynamic reward) to distinguish real offline data and generated data from the learned dynamics model. The resulting discriminator acts as a filter for out-of-distribution (or unlikely) transitions. Consequently, the paper claims through experiments that the proposed method achieves out-of-distribution generalization with the dynamic reward, therefore allowing improving existing offline MBRL methods.", "questions": "- Page 2, lines 3-4. It is unclear why the task reward function is known is reasonable given the explanation---it perhaps reinforces that reward function should not be known since it can be considered to be part of the dynamics model?\n- Page 2, last third line of first paragraph. What does it mean by \"static capabilities\"? Does it mean the model is no longer trained after (i) before equation (1)?\n\n**Possible typos**\n- Page 1, paragraph 2, line 3: \"on the models\" instead of \"on models\"\n- Page 1, paragraph 2, line 4, \"the model errors\" instead of \"model errors\"\n- In Algorithm 1, line 3: Is $\\phi$ the same as the policy $\\phi$?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The idea is simple and appears to be yield strong results in locomotion tasks.\n- The experimental analyses are thorough, in particular:\n\t- The joint distribution analysis on the learned dynamics rewards\n\t- The ablation on the filtering", "weaknesses": "The following should be addressed:\n- For the dynamics reward learning algorithm, why is GAIL better than GAN? It seems like the offline data set is fixed according to Algorithm 1, thus you can treat the dynamics models $P_t$ as the generative models while the dynamics rewards $D_t$ as the discriminators.\n- Proposition 1 is a result for tabular setting. Furthermore, this discriminator only differentiates properly the in-distribution transitions---how does this discriminator enable OOD generalization?\n\t- On page 27, last paragraph indicates that in the halfcheetah task, the dynamics reward appears to be less robust on OOD samples---is this not contradictory to what the main paper claims? Perhaps this has to do with the coverage of the data set---it appears that both `halfcheetah-medium-expert` and `halfcheetah-medium` are visibly suffering in OOD generalization, but they have less coverage on the state-action space compared to `random` and `replay`.\n- The rate is $1/\\sqrt{T}$, so how big should $T$ really be in practice? Proposition 1 somewhat suggests that $T$ should be almost covering the whole state-action space---seems like in that case we can just directly model the dynamics using a table?\n\t- Experimentally, does this correspond to 5000 as suggested by Table 4 in the appendix?\n- In experimental analysis, it appears the paper focuses mainly on locomotion tasks, and they seem to be mostly overlaps (i.e. NeoRL and D4RL share very similar environments.) I am curious if this applies to other non-trivial environments that are supported by D4RL.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)