Efficient Offline Policy Optimization with a Learned Model
Authors: Zichen Liu
Siyi Li
Wee Sun Lee
Shuicheng YAN
Zhongwen Xu
Venue: ICLR 2023 poster (2023)
Keywords: Offline RL
Model-based RL

Abstract: MuZero Unplugged presents a promising approach for offline policy learning from logged data. It conducts Monte-Carlo Tree Search (MCTS) with a learned model and leverages Reanalyze algorithm to learn purely from offline data. For good performance, MCTS requires accurate learned models and a large number of simulations, thus costing huge computing time. This paper investigates a few hypotheses where MuZero Unplugged may not work well under the offline RL settings, including 1) learning with limited data coverage; 2) learning from offline data of stochastic environments; 3) improperly parameterized models given the offline data; 4) with a low compute budget. We propose to use a regularized one-step look-ahead approach to tackle the above issues. Instead of planning with the expensive MCTS, we use the learned model to construct an advantage estimation based on a one-step rollout. Policy improvements are towards the direction that maximizes the estimated advantage with regularization of the dataset. We conduct extensive empirical studies with BSuite environments to verify the hypotheses and then run our algorithm on the RL Unplugged Atari benchmark. Experimental results show that our proposed approach achieves stable performance even with an inaccurate learned model. On the large-scale Atari benchmark, the proposed method outperforms MuZero Unplugged by 43%. Most significantly, it uses only 5.6% wall-clock time (i.e., 1 hour) compared to MuZero Unplugged (i.e., 17.8 hours) to achieve a 150% IQM normalized score with the same hardware and software stacks. Our implementation is open-sourced at https://github.com/sail-sg/rosmo.

Published as a conference paper at ICLR 2023
EFFICIENT OFFLINE POLICY OPTIMIZATION WITH A LEARNED MODEL
Zichen Liu "Sea AI Lab
Wee Sun Lee National University of Singapore {zichen, leews @ecomp nus. edu.sg
Siyi Li
Shuicheng Yan'
Zhongwen Xul
{liuze, xuzw @sea. com com
ABSTRACT
MuZero Unplugged presents a promising approach for offline policy learning from logged data. It conducts Monte-Carlo Tree Search (MCTS) with learned model and leverages Reanalyze algorithm to learn purely from offline data. For good performance, MCTS requires accurate learned models and large num- ber of simulations, thus costing huge computing time. This paper investigates few hypotheses where MuZero Unplugged may work well under the of- fline RL settings, including 1) learning with limited data coverage; 2) learning from offline data of stochastic environments; 3) improperly parameterized mod- els given the offline data; 4) with low compute budget. We propose to use regularized one-step look-ahead approach to tackle the above issues. Instead of planning with the expensive MCTS, we use the learned model to construct an advantage estimation based on one-step rollout. Policy improvements are towards the direction that maximizes the estimated advantage with regulariza- tion of the dataset. We conduct extensive empirical studies with BSuite environ- ments to verify the hypotheses and then our algorithm on the RL Unplugged Atari benchmark. Experimental results show that our proposed approach achieves stable performance even with inaccurate learned model. On the large-scale Atari benchmark, the proposed method outperforms MuZero Unplugged by 43%. Most significantly, it uses only 5.6% wall-clock time (i.e., hour) compared MuZero Unplugged (i.e., 17.8 hours) to achieve a 150% IQM normalized score with the same hardware and software stacks. Our implementation open-sourced at ltttps::/ com sai 1- sq/ rosmo
INTRODUCTION
Offline Reinforcement Learning (offline RL) (Levine et al. (020 is aimed at learning highly reward- ing policies exclusively from collected static experiences, without requiring the agent's 's interactions with the environment that may be costly even unsafe. It significantly enlarges the application potential of reinforcement learning especially in domains like robotics and health care (Haarnoja al. 2018 Gottesman al. 2019), but is very challenging. By only relying static datasets for value or policy learning, the agent in offline RI is prone to action-value over-estimation or improper extrapolation out-o--distribution (OOD) regions. Previous works (Kumar et al. 2020 Wang al., 2020 Siegel et al.) 2020) address these issues by imposing specific value penalties or policy con- straints, achieving encouraging results. Model-based reinforcement learning (MBRL) approaches
have demonstrated effectiveness in offline RL problems (Kidambi et al. 2020 Yu et al. 2020 Schrit- twieser et al. 2021). By modeling dynamics and planning. MBRL learns as much as possible from the data, and is generally more data-efficient than the model-free methods. We are especially inter- ested in the state-of-the-art MBRL algorithm for offline RI, i.e., MuZero Unplugged (Schrittwieser et al. 2021), which is a simple extension of its online RL predecessor MuZero (Schrittwieser al. 2020). :. MuZero Unplugged learns the dynamics and conducts Monte-Carie Tree Search (MCTS) (Coulom 2006, Kocsis & Szepesvári 2006) planning with the learned model to improve the value and policy in a fully offline setting.
In this work, first scrutinize the MuZero Unplugged algorithm by empirically validating hy- potheses about when and how the MuZero Unplugged algorithm could fail in offline RL settings.
"Corresponding author.

---

Published as a conference paper at ICLR 2023
The failures could also happen in online RL settings, but the intrinsic properties of offline RL mag- nify the effects. MCTS requires an accurate learned model to produce improved learning targets. However, in offline RL settings, learning accurate model is inherently difficult especially when the data coverage is low. MuZero Unplugged is also not suitable to plan action sequences in stochas- tic environments. Moreover, MuZero Unplugged is compute-intensiig algorithm that leverages the compute power of an NVIDIA V100 x One week for running each Atari game (Schrittwieser et al.) 2021]. When trying to reduce the compute cost by limiting the search, MuZero Unplugged fails to learn when the number of simulations in tree search is low. Last but not least, the implementation of MuZero Unplugged is sophisticated and dlose-sourced, hampering its wide adoption in the research community and practitioners.
Based on the hypotheses and desiderata of MBRL algorithms for offline RI., we design ROSMO, a Regularized One-Step Model-based algorithm for Offline reinforcement learning. Instead of conducting sophisticated planning like MCTS, ROSMO performs a simple yet effective one-step look-ahead with the learned model to construct improved target for policy learning and acting. To avoid the policy being updated with the uncovered regions of the offline dataset, we impose policy regularization based on the dataset transitions. We confirm the effectiveness of ROSMO first on BSuite environments by extensive experiments on the proposed hypotheses, demonstrating that ROSMO is more robust model inaccuracy, poor data coverage, and learning with data of stochastic environments. We then compare ROSMO with state-of-the-an methods such as MuZero Unplugged (Schritwieser et al.] 2.2I). Critic Regularized Regression (Wang et al., 2020). Conser- vative Q-Learning (Kumar et al. 2020). and the vanilla Behavior Cloning baseline in both BSuite environments and the large-scale RL Unplugged Atari benchmark (Gulcehre et al. 2020). In the Atari benchmark, the ROSMO agent achieves a 194% IQM normalized score compared to 151% of MuZero Unplugged. It achieves this within fraction of time (i.e., hour) compared MuZero Unplugged (i.e.. 17.8 hours), showing improvement of above 17x in wall-clock time, with the same hardware and software stacks. We conclude that a high-performing and easy-to-understand MBRL algorithm for offline RL problems is feasible with low compute resources.
BACKGROUND
2.1 NOTATION
The RL problem is typically formulated with Markov Decision Process (MDP), represented by M P,r,,, with the state and action spaces denoted by S and A, the Markovian transition dynamics P. the bounded reward function the discount factor y and an initial state distribution p. At any time step, the RL agent in some state s E S interacts with the MDP by executing an action E A according The value function of a fixed policy and a starting state 8p = 8 ~ p is defined as the expected cumulative discounted rewards (*(() E |E Ee= 7'r (st, so s). An offline RL agent aims to learn a policy æ that maximizes (s))) = (*()) solely based on the static dataset D, which contains the interaction trajectories {T.} of one or more behavior policies Ty with M. The learning agent π cannot have further interaction with the environment to collect more experiences.
policy (a]s), arrives at the next state s' and obtains reward r(s, a, s') E R.
2.2 OFFLINE POLICY IMPROVEMENT VIA LATENT DYNAMICS MODEL
Model-based RL methods (utton 1991
promising to solve the offline RL problem because they can effectively learn with more supervision signals from the limited static datasets. Among them, MuZero learns and plans with the latent model to achieve strong results in various domains (Schrittwieser et al. 2020). We next describe a general algorithmic framework extended from MuZero for the offline RL settings, which we refer as Latent Dynamics Model.
Deisenroth & Rasmussen 2011
Given a trajectory = (ot, 41, fl...., OT,, OT,, E D and at any time step I. E T.). we encode the observations into a latent state via the representation function hg: s = he (ot). We could then unroll the learned model recurrently using the dynamics function ge to obtain an imag- ined next state in the latent space and an estimated reward: ,+1 *+1
(s/ ar+.), where k E [0, K] denotes the imagination depth (the number of steps we unroll using the learned model go). Conditioned on the latent state the prediction function estimates the policy and value function:
2

---

Published as a conference paper at ICLR 2023
T of = fo (st). Note that we have two timescales here. The subscript denotes the time step in trajectory, and the superscript denotes the unroll time step with the learned model.
In the learning phase, the neural network parameters updated via gradient descent over the loss function as follows.
K ((( = Σ Ex (1t, (4tk) + E" (vf, + ET (o(:PPt)) + t00 (4).
(1)
O=
where E", td are loss functions for reward, value and policy respectively, and fae can be any form of regularizer. The exact implementation for loss functions can be found in Appendii A.2, "244 is the reward target from the environment (i.e., dataset), Pt+kkPtt = Dg ($t+*) are the value and policy targets output from an improvement operator I using dynamics and prediction functions with target network parameters 6' (Mnih et al. (013). Note that the predicted next state from the dynamics function is not supervised for reconstruction back to the input space. Instead, the model is trained implicitly so that policy and value predictions at time-step from imagined states at depth & can match the improvement targets real time-step I + from the environment. This is an instance of algorithms that apply the value equivalence principle (Grimm et al. 2020). Improvement operators can have various forms; for example, MuZero Unplugged (Schrittwieser et al. 2021 applies Monte-Carle Tree Search.
2.3 MONTE-CARLO TREE SEARCH FOR POLICY IMPROVEMENT
We briefly revisit Monte-Carlo Tree Search (MCTS) (Coulom) 2006 Kocsis & Szepesvári 2006), which can serve as an improvement operator to obtain value and policy targets. To compute the targets for o. and u/. start from the root state st44. and conduct MCTS simulations up to a budget N. Each simulation traverses the search tree by selecting actions using the pUCT rule Rosin 2011):
no,. 1 7 n(s
(cs (a log Es [os + (Eom(s,))+aatt)
a² arg max Q(s, a) + prii (8, a).
C2
(2)
where n(s, a) is the number of times the state-action pair has been visited during search, Q(s.a) is the current estimate of the Q-value, Tpior((, a) is the probability of selecting action a in state s using the prior policy, and C1, C2 are constants. When the search reaches a leaf node s will expand the tree by unrolling the learned model go with a' and appending new node with model predictions vi I+1 1+1
me [+1 _11 _zt11 to the search tree. Then the estimate of bootstrapped discounted
r- y5 -T,4+1+ + - for = l... is backed up all the way to the root node, updating Q and statistics along the path. After exhausting the simulation budget, the policy target PL++ is formed by the normalized visit counts at the root node, and the value target Z4+8 is the n-step discounted return bootstrapped from the estimated value at the root node
return G 1
ni(H n((8,a)2/" a)2/T
PMCTs(al(1)
[ns 6)1/T
(3)
n n(bb-n;a)) a)
4+-11
ZMCTS($:) >
Q( Q(sttn: +
Σ
env
EB ( n(notn)bb))
t=
METHODOLOGY
3.1
MOTIVATION
With well learned function estimators (ho, fe, 3o planning with MCTS (Scctinn 2.3 in the latent dynamics model (Section 2.2) has been shown to obtain strong policy and value improvements, and has been applied to offline RL (Schrittwiser prohibitive computational power required by the search is limiting its practicability. For example,
al. 2021) in MuZero Unplugged. However, the
We re-index t t & for an easier notation.
3

---

Published as a conference paper at ICLR 2023
Figure 1: An illustration on the differences of Monte-Carlo Tree Search (left) and one-step look-ahead (right) for policy improvement. With limited offline data coverage, the search entering regions bcyond the observed data may expand nodes on which fo fails to provide accurate estimations, and ge even leads to worse region. These errors are compounded along the path as the search gets deeper. leading to detrimental improvement largets. One-step look-ahead is less likely to go outside the safe region. Illustrations are best viewed on screen.
experiments the Atari benchmark (Bellemare al., 2013) could take NVIDIA V100 One week to run a single game. Besides, in the offline RL settings, the state-action space coverage of the dataset is inherently limited, and further environment exploration is not possible. Thus, learned estimators may only be accurate in the safe regions covered by the dataset, and generalization outside the safe regions may lead to extrapolation error (Fujimoto al. 2019.. Even worse, the MCTS process unrolls the learned model recurrently with actions selected using Equation (which depends on the estimation of Tprrrr), compounding the extrapolation errors along the search path.
Figure T leftt illustrates when MCTS goes wrong as an improvement oper- ator. Suppose we have encoded the observations through hi to get latent states 8t, 8 8142, which we refer to as observed nodes, as shown in the figure. The dotted edges and nodes in the left figure describe a sim- ple MCTS search tree at node 8t, with simulation budget N = 7 and depth 3. 3. We refer to the states unrolled by the learned model, {s{] J. as imag- ined nodes, with 1,2, denoting the depth of search. The regions far from the observed nodes and beyond the shaded area are considered unsafe regions, where few no data points were collected. Intuitively, policy and value predictions on imagined nodes in unsafe regions are likely erroneous, and so are the reward and next state imaginations. This makes the improvement targets obtained by Equation 3 Junreliable. Furthermore, we also argue that MuZero Unplugged lacks proper regularization that constrains the policy from going too far away from the behavior policy to combat the distributional shift. In their algorithm, ged = c|l0| 22 Equation 1) only imposes weight decay for learning stability.
Algorithm ROSMO
Require: dataset D, initialized model parameters
I: while True do
2: Sample a batch of trajectory B E
3:
2 f- compute value target (Equation 8
4:
Pr + compute policy target (Equation 6
5:
grog <- apply regularization (Equation 11
6:
l f- compute loss (Equation I) on B
Update 9 with gradient descent on
8: end while
Motivated by the analysis above, the desiderata of a model-based offline RL algorithm are: com- pute efficiency, robustness to compounding extrapolation errors, and policy constraint. To addresses these desiderata, we design ROSMO, a Regularized One-Step Model-based algorithm for Offline reinforcement learning based on value equivalence principle. As illustrated in Figure righht, ROSMO performs one-step look-ahead to seek the improvement targets from a learned model, which more efficient than MCTS and less affected by compounding errors. The overview of our complete algorithm is described in Algorithm I] The algorithm follows [Section 2.2 to encode the observations, unrolls the dynamics, makes predictions, and computes the loss. The blue triangles highlight our algorithmic designs on learning targets and regularization loss. We will derive them in the following sections.
3.2 A SIMPLE AND EFFICIENT IMPROVEMENT OPERATOR
In the algorithmic framework we outlined in [Section 2.2, the improvement operator is used to compute improved policy and value target. Unfortunately, MuZero Unplugged uses the compute-
4

---

Published as a conference paper at ICLR 2023
heavy and sophisticated MCTS to achieve this purpose, which is also prone to compounding trapolation errors. In particular. the policy update of MuZero Unplugged is done by minimizing the cross entropy between the normalized visit counts and the parametric policy distribution:
n(s, a) Es n/s,,()/77 b)1/T logm(als))
MCTS = E
(4)
where the visit counts n at the root node s are summarized from the MCTS statistics.
We propose learn the value equivalent model and use a more straightforward and much more efficient one-step look-ahead method to provide policy improvement. Specifically, our policy up- date is towards minimizing the cross entropy between a one-step (OS) improvement target and the parametric policy distribution:
lss -p-pppa log π
(5)
The policy target p state is estimated as:
p(als) = Tprior (als) exp (adv, (8,
(6)
(())
where prior the prior policy (often realized by the target network ), adv, (s, a) - Og (s a)
(s) is approximation of the action advantage from the learned model ge. and the factor (s) ensures the policy target is a valid probability distribution. The state value (s)) = (s) is from the current state. The action value ag (s, a) is estimated by
the prediction function conditioned unrolling the learned model one step into the future using dynamics function go to predict the reward rg and next state s and then estimate the value at the imagined next state:
Tgx S'g go(s. a), (s, a) = rg + f5o (syy
Intuitively, our policy target from Equation jjdjuusss the prior policy such that actions with positive advantages are favored, and those with negative advantages are discouraged.
Mcanwhile, the value target is the n-step return bootstrapped from the value estimation:
(st) 0r4n
Σ
rre
(8)
where Uetn fer. (St+n) is computed using the target network l' and re- is from the dataset. Compared to MuZero Unplugged, the value target is simpler, eliminating the dependency on the searched value the node n steps apart.
Sampled policy improvement. Computing the exact policy loss l needs to simulate all actions to obtain full p(a|s) distribution, and then apply the cross entropy. It demands heavy computation for environments with large action space. We thus sample the policy improvement by computing an estimate of S on N < |A| actions sampled from the prior policy, i.e., a19) (s):
N
exp (adv a()))
los
log (a))(s)))))
(9)
N
Z(s)
The normalization factor ((s) for the k-th sample out of N can be estimated (Hessel et al. 2021
1 Eiyx exp (adv ()))
2 (s)
(10)
N
In this way, the policy update will not grow its computational cost along with the size of the action space.
3.3 BEHAVIOR REGULARIZATION FOR POLICY CONSTRAINT
Although the proposed policy improvement in Section 3.2 alleviates compounding errors by only taking one-step look-ahead, could still be problematic if this step of model rollout leads
5

---

Published as a conference paper at ICLR 2023
to imagined state beyond the appropriate dataset extrapolation. Therefore, some form of policy constraint is desired to encourage the learned policy to stay close to the behavior policy.
To this end, we extend the regularization loss to S (0) c| be any regularization applied to reward, value, or policy predictions and is jointly minimized with other losses. While sophisticated regularizers such penalizing out-of-distribution actions' reward predictions could be designed, we present a simple behavior regularization on top of the policy म, leaving other possibilities for future research.
The second term can
Our behavior regularization is similar to Siegel et al. (2020), but we do not learn an explicit behavior policy. Instead, we apply an advantage filtered regression directly on from the prediction function output:
F 2 = (ssa)>> |- log (als)) H (adv (s),)))
(11)
where H(x) = 1 2 is the Heaviside step function. We can interpret this regularization objective as behavior cloning (maximizing the log probability) on a set state-action pairs with high quality (advantage filtering).
3.4 SUMMARY
Our method presented above unrolls the learned model for one step to look ahead for improve- ment direction and adjusts the current policy towards the improvement with behavior regularization. As illustrated in Figure compared with MuZero Unplugged, safe regions with higher probability and utilize the appropriate generalization to improve the policy. Our policy update can also be interpreted as approximately solving a regularized policy optimization problem, and the analysis can be found in Appendix B
method could stay within the
4 EXPERIMENT
In Figure we have analyzed the deficiencies of MuZero Unplugged and introduced ROSMO as simpler method to tackle the offline RL problem. In this section, we present empirical results to demonstrate the effectiveness and efficiency of our proposed algorithm. We firstly focus on the comparative analysis of ROSMO and MuZero Unplugged to justify our algorithm designs in |Sec- fion 4.1, and then we compare ROSMO with existing offline RL methods and ablate our method in Section 4.2) Throughout this section we adopt the Interquartile Mean (IQM) metric (Agarwal al. 2021 on the normalized scord to report the performance unless otherwise stated.
4.1 HYPOTHESIS VERIFICATION
To analyze the advantages of ROSMO over MuZero Unplugged, we put forward four hypotheses for investigation (listed in bold) and verify them on the BSuite benchmark (Osband et al.) 2019).
Similar to (Gulcehre et al. 2021). we use three environments from BSuite: catch, cartpole and mountain_car, However, for training model-based agents since it only contains unordered transitions instead of trajectories. Therefore, we generate an episodic dataset by recording experiences during the online agent training and use it in our experiments (see Appendixx E for data collection details).
note that the released dataset by Gulcehre et al. (202T) is unsuitable
(1) MuZero Unplugged fails to perform well in low-data regime. With low data budget, the coverage of the state-action space shrinks as well as the safe regions (Figure 1}. We hypothesize that MuZero Unplugged fails to perform well in the low-data regime since the MCTS could easily enter the unsafe regions and produce detrimental improvement targets. FFiuree 2/a) shows the IQM normalized score obtained by agents trained on sub-sampled datasets of different fractions. We can observe MuZero Unplugged degrades when the coverage becomes low and performs poorly with 1% fraction. In comparison, ROSMO works remarkably better in a low-data regime and outperforms MuZero Unplugged across all data coverage settings.
Calculated as (score scorCrandom) )/ (scorcualine score per game, = 1 means par with the online agent used for data collection; x indicates its performance is 1.5x of the data collection agent's.

---

Published as a conference paper at ICLR 2023
MZU
ROSMO
2,034.01 4101
0.980
1.0.c0
0.1
0.984
1.0aw
0.3
0.772
0.900
AD9PD
0.36
0.5
0.404
0.916
15 40 25
tumns percetaao
Figure 2: (a) IQM normalized score with different data coverage. (b) IQM normalized score with ditferent dynamics model capacities. (c) IQM normalized score of MuZero Unplugged with different simulation budgets (N) and search depths (d).
Table 1:
Comparison between ROSMO and MuZero Un- plugged on catch with different noise levels (c).
(2) ROSMO is more robust in learning from stochastic transitions than MuZero Unplugged. To evaluate the robustness of MuZero Unplugged and ROSMO in learning with data from stochastic environments, we inject noises during experience collection by replacing the agent's action with a random action for environment execution, with probability stochastic, MuZero Unplugged could fail to plan action sequences due to compounding errors. We hypothesize that ROSMO performs more robustly than MuZero Unplugged since ROSMO only uses a one-step look-ahead, thus has less compounding error. In Table 1, we compare the episode return of the two algorithms with the controlled noise level. The result shows that ROSMO is much less sensitive to the dataset noise and can learn robustly at different stochasticity levels.
With the dataset dynamics being
(3) MuZero Unplugged suffers from dynamics mis- parameterization while ROSMO is less af- fected. The parameterization the dynamics model crucial for model-based algorithms. It is difficult to design model with the expressive power that is appropriate for learning the dataset's MDP transition. The resulting under/over fitting of the learned model may badly affect the per- formance of the overall algorithm. We hypothesize that MuZero Unplugged is more sensitive to the parameterization of dynamics than ROSMO. Figure 2(b) compares ROSMO with MuZero Un- plugged for different dynamics model capacities trained on 10% data. Since we use multi-layer perceptron to model the dynamics function, the capacity is controlled by the number of hidden units. We show that MuZero Unplugged works best when the number of hidden units is 1024, and its per- formance degrades significantly with less model capacity, likely due to the under-fitting of smaller networks. The effect of over-fitting less obvious. In comparison, ROSMO performs stably with different dynamics model capacities and consistently outperform MuZero Unplugged in all settings.
(4) MuZero Unplugged is sensitive to simulation budget and search depth. Prior works have shown that the performance of MuZero agents declines with a decreasing simulation budget (Grill cct al. 2020), and it is insensitive to search depth (Hamrick et al. 2021). Both works consider online RL settings, where new experience collection may correct prior wrong estimations. We hypothesize that in offline RL settings, the performance of MuZero Unplugged is sensitive to both simulation budget and search depth. In particular, deeper search would compound extrapolation errors in offline settings, leading to harmful improvement targets. Figure 2[c) demonstrates the IQM nor- malized score of MuZero Unplugged with different simulation budgets and search depths. We can observe MuZero Unplugged fails to learn when N is low, and it performs poorly when N is high but with a deep search. This suggests that too low visit counts are not expressive, and too much planning may harm the performance, matching the findings in the online settings (Grill et al. [2020] Hamrick et al., (02T). Notably, limiting the search depth can ease the issue by large amount, serving as further strong empirical evidence to support our hypothesis that deep search compounds errors, reinforcing our belief in the one-step look-ahead approach.
4.2 BENCHMARK RESULTS
After investigating what could go wrong with MuZero Unplugged and validating our hypotheses, we compare our method with other offline RL baselines on the BSuite benchmark as well as the larger-scale Atari benchmark with the RL Unplugged (Gulcehre al. 2020) dataset.
Baselines. Behavior Cloning learns a maximum likelihood estimation of the policy mapping from the state space to the action space based the observed data, disregarding the reward signal. Thus BC describes the average quality of the trajectories in the dataset and serves as a naive baseline. Conservative Q-Learning (CQL.) (Kumar al. 2020 learns lower-bounded action values by incor-
7

---

Published as a conference paper at ICLR 2023
porating loss penalties the values of uut-of-distributtin actions. Critic Regularized Regression (CRR) (Wang et al. 2020) approaches offline RL in a supervised learning paradigm and reweighs the behavior cloning loss via an advantage estimation from learned action values. CQL and CRR representative offline RL algorithms with strong performance. MuZero Unplugged (MZU) (Schrit-) twieser et al.] 202T) is a model-basee method that utilizes MCTS plan for learning well as acting, and exhibits state-of-the-art performance on the RL Unplugged benchmark. MOReL (Ki-) dambi al. 2020 and MOPO (Yu et al.. 2020) are another two model-based offline RL algorithms. MOReL proposes to learn a pessimistic MDP and then learn a policy within the learned MDP; MOPO models the dynamies uncertainty to penalize the reward of MDP and learns a policy on the MDP. Both of them focus on state-based control tasks and are not trivial to transfer to the image- based Atari tasks, hence are not compared here. Nevertheless, we do provide the details of our implementation for MOReL. (Kidambi et al. 2020) and COMBO Yuu et al. 202T) (which extends MOPO (Yu et al. 2020)) and report the results in Appendix F.I.
Implementation. We use the same neural network architecture to implement all the algorithms and the same hardware to run all the experiments for a fair comparison. We closely follow MuZero Unplugged (Schrittwieser et al.) 2021) to implement ROSMO and MuZero Unplugged, but use down-scaled version for Atari to trade off the experimentation cost. We use the exact policy loss without sampling (Equation 5) for all ROSMO experiments in the main results and compare the performance of sampling in the ablation. For CRR and CQL we adapt the official codes for our experiments. For Atari we conduct experiments on set of 12 Atari games due to limited computa- tion resources We ensure they are representative and span the performance spectrum of MuZero Unplugged for fair comparison. More implementation details can be found in Appendix D.2
normalized score
2.0
Method
MountainCar
Canpole
1.3
BC
0.
24-4112
589
ROSMO
CQI.
1.0gga
124.77 77
416.48
1.0
MZU
CRR
1. 1.0gea
-006.0 3at
907. 1o
MZU.
0. a.
-107. 2. es
890.64217h
0.5
CRR
ROSMO
1.0gea
080.68ate
CQL
BC
0.0
100
20
Steps (in thousands)
Time (in hours)
Table 2: BSuite benchmark results. Averaged episode returns measured at the end of training (200K steps) across 5 seeds.
Figure 3: Atari benchmark results. Aggregated IQM nor- malized score of different algorithms in terms of (left) sample efficiency and (right) wall-clock efficiency.
Main results. Table 2] shows the BSuite benchmark results of our algorithm and other baseline methods. ROSMO achieves the highest episode returns on catch and mountain with the lowest standard deviation. For cartpole, CRR performs slightly better than ROSMO, but we still observe ours outperforms the other baseline by clear margin.
Figure leftt) presents the learning curves with respect to the learner update steps, where it is clearly shown that ROSMO outperforms all the baselines and achieves the best learning efficiency. In terms of the final IQM normalized score, all offline RL algorithms outperform the behavior cloning baseline, suggesting that the reward signal is greatly helpful when the dataset contains trajectories with diverse quality. ROSMO obtains a 194% final IQM normalized score, outperforming MuZero Unplugged (151%), Critic Regularized Regression (105%), and Conservative Q-Learning (83.5%). Besides the highest final performance, we also note the ROSMO learns the most efficiently among all compared methods.
Figure 3 right) compares the wall-clock efficiency of different algorithms given the same hard- ware resources. Significantly, ROSMO uses only 5.6% wall-clock time compared to MuZero Un- plugged to achieve 150% 1QM normalized score. With the lightweight one-step look-ahead design, the model-based ROSMO consumes similar learning time as model-free methods, widening its ap- plicability both offline RL researches and real-worid applications.
O 'Even with our slimmed implementation (Appendix for details), full run of all compared methods on the RI. Unplugged benchmark (46 games) needs about 720 TPU-days for a single seed, which is approximately equivalent NVIDIA V100 GPU running for 2 years.
8

---

Published as a conference paper at ICLR 2023
PGL
Narder oi In thopenbl
Numbor ototes In thossands
Number Japs On teesandd)
Mamber Stop On despandi)
(8)
(1)
(4)
Figure 4: Learning curves of IQM normalized score on MsPacman. (a) Comparison of ROSMO and MuZero
Unplugged in low data regime. (b) Comparison of ROSMO, MuZero Unplugged and MZU-Q when limiting the number of simulations (number of samples) be N 4. (c) Comparison of ROSMO and MuZero Unplugged when the model is unrolled with different steps for learning. (d) Ablation of the one-step policy improvement and the behavior regularization.
The results of individual games can be found in Appendix F
Ablations. We present our ablation studies on data coverage, learning efficiency, model compound-
ing errors, and decoupled ROSMO. Following the common practice (Schrittwieser al. 2020 Hamrick et al. 2021), Ms. Pacman is chosen for the ablation studies.
(a) Figure 4 aa) shows that ROSMO is able to outperform MuZero Unplugged in both 10% and 1% data regimes, replicating our hypothesis verification results on BSuite.
(b) To make MuZero Unplugged more compute-cfficient and feasible, could limit the number of simulations. However, prior works have shown that MuZero's 's policy target degenerates under low visit count (Grill et al. 2020 Hamrick et al. (020)_ Hence, we also implement the MZU-Q variant which uses an MPO-style (Abdolmalcki et al. 2018) policy update, π xMPO x ли exp(() MQMCTS /T), for comprchensive comparison. Here Que
the Q-values at the root node of the search tree, and r is a temperature parameter set to be 0.1 following Hamrick et al.] (202T). Figure 4(b) shows that MZU fails to learn using 4 simulations, while MZU-Q can somewhat alleviate the issues. Our sampled ROSMO performs well with a limited sampling budget.
(c) To alleviate the compounding errors (Janerr et al. 2019), MuZero Unplugged unrolls the dy- namics for multiple steps (5) and learns the policy, value, and reward predictions on the recurrently imagined latent state to match the real trajectory's 's improvement targets. It also involves compli- cated heuristics such scaling the gradients at each unroll step to make the learning more stable. We ablate ROSMO and MuZero Unplugged using a single-step unroll for learning the model. Fig- ure 4[c) shows that the performance of ROSMO is not sensitive to the number of unrolling (either or 5). while MuZero Unplugged experiences a significant performance drop when only single-step unrolling is applied.
(d) We ablate the effect of behavior regularization. We compare our full algorithm ROSMO with two variants: OneStep one-step policy improvement without regularization; Behavior policy learning via standalone behavior regularization. Interestingly, the Behavior variant recovers the binary form of CRR (Wang et al.] 2020). with an advantage estimated from the learned model. Figure 4dd) demonstrates that compared to OneStep. Behavior learns more efficiently at the early stage by mimicking the filtered behaviors of good quality, but is gradually saturated and surpassed by OneStep, which employs a more informative one-step look-ahead improvement target. However, OneStep alone without behavior regularization is not enough especially for low-data regime (see Appendix F.3 for more results). The full algorithm ROSMO combines the advantages from both parts to achieve the best learning results.
CONCLUSION
Starting from the analysis of MuZero Unplugged, we identified its deficiencies and hypothesized when and how the algorithm could fail. We then propose our method, ROSMO, a regularized one- step model-based algorithm for offline reinforcement learning. Compared MuZero Unplugged, the algorithmic advantages of ROSMO are threefold: (1) it is computationally efficient, (2) it is robust compounding extrapolation errors, and (3) it is appropriately regularized. The empirical investigation verified our hypotheses and the benchmark results demonstrate that our proposed al- gorithm can achieve state-of-the-un results with low experimentation cost. We hope our work will serve as a powerful and reproducible agent and motivate further research in model- based offline reinforcement learning.

---

Published as a conference paper at ICLR 2023
ETHICS STATEMENT
This paper does not raise any ethical concerns. Our study does not involve human subjects. The datasets we collected do contain any sensitive information and will be released. There are no potentially harmful insights or methodologies in this work.
REPRODUCIBILITY STATEMENT
To ensure the reproducibility of our experimental results, we included detailed pseudocode, im- plementation specifics, and the dataset collection procedure in the Appendix. More importantly, we released our codes as well as collected datasets for the research community to use, adapt and improve our method.
REFERENCES
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Mar- tin Riedmiller. Maximum posteriori policy optimisation. In ICLR, 2018.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. NeurIPS. 34:29304-29320 2021.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The Arcade Learning Envi- ronment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Minn and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http:
//github. com/ google/jaaaeaaoooooo
Rémi Coulom. Efficient selectivity and backup operators Monte-Carlo Tree Search. Interna- tional conference computers and games, pp. 72-83. Springer, 2006.
Marc Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-1I), pp- 465-472. Citeseer, 2011 1.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In ICML. pp. 2052-2062. PMLR, 2019.
Omer Gottesman, Fredrik Johansson, Matthieu Komorowski, Aldo Faisal, David Sontag, Finale Doshi-Veles Velez, and Leo Anthony Celi, Guidelines for reinforcement learning healthcare. Nature medicine, 25(1):16-188 2019.
Jean-Bastien Grill, Florent Altché, Yunhao Tang, Thomas Hubert, Michal Valko, Ioannis
Antonoglou, and Remi Munos. Monte-Carlo Tree Search regularized policy optimization. In ICML, volume 119, pp. 3769-3778. PMLR, 2020.
Christopher Grimm, André Barreto, Satinder Singh, and David Silver. The value equivalence prin- ciple for model-basee reinforcement learning. NeurIPS, 33:5541-5552, 2020.
Caglar Gulcehre, Ziyu Wang. Alexander Novikov, Thomas Paine, Sergio Gomez Colmenarejo, Kon- rad Zolna, Rishabh Agarwal, Josh Merel, Daniel J. Mankowitz, Cosmin Paduraru, Gabriel Dulac- Arnold, Jerry Li, Mohammad Norouzi, Matthew Hoffman, Nicolas Heess, and Nando de Freitas. RI unplugged: A collection of benchmarks for offline reinforcement learning. In NeurIPS, 2020.
10

---

Published as a conference paper at ICLR 2023
Caglar Gulcehre, Sergio Gómez Colmenarejo, Ziyu Wang, Jakub Sygnowski, Thomas Paine, Konrad Zolna, Yutian Chen, Matthew Hoffman, Razvan Pascanu, and Nando de Freitas. Regularized behavior value estimation. arXiv preprint arXiv:2103.09575. 2021.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Schoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and appli- cations. arXiv preprint arXiv:1812.05940 2018.
Danijar Hafner, Timothy P Lillicrap. Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In ICLR, 2021.
Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez Tobias Pfaff, Theophane Weber, Lars Buesing, and Peter W. Battaglia. Combining q-learning and search with amortized value esti- mates. In ICLR, 2020.
Jessica B Hamrick, Abram I.. Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Wither- spoon, Thomas Anthony, Lars Holger Buesing, Petar Veličković, and Theophane Weber. On the role of planning in model-based deep reinforcement learning. In ICLR, 2021.
Matteo Hessel, Ivo Danihelka. Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre. Theophane Weber, David Silver, and Hado Van Hasselt. Muesli: Combining improvements in policy opti- mization. In ICML., volume 139, pp. 4214 4226. PMLR, 2021.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. NeurlPS, 32, 2019.
Sham M. Kakade and John Langford. Approximately optimal approximate reinforcement learning. In ICML, pp. 267-274, 2002.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. MOReL: Model- based offline reinforcement learning. NeurIPS, 33:21810-21883 2020.
Levente Koesis and Csaba Szepesvári. Bandit based Monte-Carke planning. ECMI,, pp. 282-293. Springer, 2006.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. NeurIPS, 33:1 13:1119-1 2020.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01644; 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing Atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvari, Satinder Singh, et al. Behaviour suite for reinforce- ment learning. arXiv preprint arXiv: 100Xiv:1000.000 03566, 2019.
Tobias Pohlen, Bilal Piot, Todd Hester, Mohammad Gheshlaghi Azar, Dan Horgan, David Budden, Gabriel Martn, Hado Van Hasselt, John Quan, Mel Večerík, et al. Observe and look further: Achieving consistent performance on atari. arXiv preprint arXiv: 1593, 2018.
Christopher D Rosin. Multi-armed bandits with episode context. Annals of Mathematics and Artifi- cial Intelligence, 633) 2011.
Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, and David Silver. Mastering Atari, Go, Chess and Shogi by Planning with Learned Model. Nature, 588(7839):604-600)) 2020.
Julian Schrittwieser, Thomas Hubert. Amol Mandhane, Mohammadamin Barckatain, Ioannis Antonoglou, and David Silver. Online and offline reinforcement learning by planning with a learned model. In NewrlPS, volume 34, pp. 27580-27591, 2021.

---

Published as a conference paper at ICLR 2023
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In ICML, pp. 1889-1897. PMLR, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint vrXiiv:107.063444 2017.
Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Ne- unert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavioral modelling priors for offline reinforcement learning. arXiv preprint arXiv:2002.0839% 2020.
Richard S Sutton.
Dyna, integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 164) 1991.
Richard S Sutton and Andrew G Barto. Introduction to reinforcement learning. 2020.
Cameron Voloshin, Hoang Le, Nan Jiang, and Yisong Yue. Empirical study of off-policy policy evaluation for reinforcement learning. NeurIPS Track on Datasets and Benchmarks, volume 1, 2021.
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. NeurIPS, 0167768-7777 2020.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: Model-based offline policy optimization. NerulPS, 33:14129-14142, 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. In NeurlPS, volume 34, pp. 28954-28967, 2021.
12

---

Published as a conference paper at ICLR 2023
A ALGORITHMIC DETAILS
A.I PSEUDOCODE
We present the detailed learning procedure of ROSMO in Algorithm 2 we use a single slice of trajectory, but in practice we can take batches for parallel.
For notational convenience,
Algorithm 2 ROSMO Pscudocode
Require: dataset D, discount factor 7, initialized model parameters 0, target parameters 0' 8, unroll step K, TD step n, behavior regularization strenth a, weight decay strength sampling budget N (optional)
1: while True do
2: Sample a trajectory Ti E D with length T;
3: Sample random time step / E [0, T; K
4: of <- ho(
representation of root
5: Tr. Se, We. V2 < UNROLL( si, Ot.....+K-1)
6:
Pr. Zt f IMPROVE(O', ot.....+KKt 1 t,...t- K+n)
7:
greg < RBEHAVORRREULLRRZZR (0) Te, Ot.....-K ag (K) + c|01
8:
rew <- directly get from тr indexed at 1,..., K
9: E",e", "" +- compute losses following Appendix A.2
Update with gradient descent on f" Ed + e update O' 0 with interval end while
11:
12: function UNROLL( 50 (u.)...K11))
13: initialize vector containers r, n,, with r" = 0,8 = 0
14: for j = O...K K do
15:
n v' t fo(s") rtt s/+1 1't1,s/11
prediction on root and the imaginary dynamics
go(s', a,)
17: end for
TK v < k)
19:
return s, v
20:
end function
21:
function IMPROVE(8', Oo.... +ni ?,... K.n...K(n.)
initialize vector containers p, z for policy and value targets
23:
hg (00...,KKnn)
24:
n, v + (s) for O...K do
25:
26:
adv +- ONESTEPLOOKAHHAA(((
vi)
27:
p' т exreeppddd/////v
z < "vntts + Er= 1utn-1 g6 -t
end for
30: return p.
31:
end function
32: function ONESTEPPOOKAHEADDEE
33:
f- sample N enumerate actions
34: r,s' f- 90 (s,
35: return r "
36: end function
37:
function BEHAVIORREGULARIZZERR her (ot....KK1))
Oo....,K Co.....K K)
39: r.s' +- g0 (s, ao....
adv <- ((00'((s') (0)
41: return सनन Ex= log( T
42: end function
13

---

Published as a conference paper at ICLR 2023
A.2 TRAINING
An illustration of the training procedure can be found in Figure Seleft). To optimize the network weights, we apply gradient descent updates on 0 over the loss defined in Equation Specifically, our loss functions for policy, value and reward predictions are:
(T p) -pp 8" (v, " log ( (r, = logr,
T log ,
(12)
(13)
(14)
where z' h(=), v' h(r=m)) are the value and reward targets scaled by the invertible transform h(x) (ign(x)( 1/ ez). where 0.001 (Pohlen al. 2018). We then apply transformation d to obtain the equivalent categorical representations of scalars, which then serve as
the targets of cross entropy loss of the scalars' distribution predictions. For the policy prediction, the loss is its cross entropy with the improved policy target.
In Algorithm 2 (line-9), we have vectorized inputs of length K + for the loss computation, thus for every element we apply the above loss functions and take the average.
We also follow MuZero (Schrittwieser et al. 2020) closely to scale the gradients at the start of dynamics function. However, we do not use the prioritized replay for simplicity and do not apply the normalization to the hidden states.
B ANALYSIS
We can interpret minimizing the policy loss in Equation 12 as conducting a regularized policy op- timization. Suppose our goal is to maximize the expected improvement 17(x)) = J J over the behavior policy p(als) = Tg(a)s). which can be expressed terms of the advantage adv a)
with respect to y (Kakade & Langford 2002 Schulman et al. 2015):
nTT)) Egnd, (1) (s)Ba= (s.
(15)
where d, = Ei P(s s|x) is the unnormalized discounted distribution of state visitation induced by policy и (Sutton & Barto 2020). practice, we follow [Schulman et al.) (2015) to optimize an approximation f)(x)) = Esndeee Bawa((a)) s) (s, a)], which provides a good estimate of TT(T)) when and и are close (S e) in terms of KL-divergence. Hence, can use as the surrogate objective and maximize under a constraint, serving as a regularized policy optimization:
/. r(ass) (s, a)] dads J. du((s)
arg max
(16)
s.t. [ (4L6))Daa( s..
Using the Lagrangian of [qqatton 16 the optimal policy can be solved as:
1 (als) Z(s) p(als) exp (adv (s, a))
(17)
where 7(s) is the partition function, and B is a Lagrange multiplier. The learning policy can be improved via projecting back onto the manifold of parametric policies by minimizing their KL- divergence:
arg min EgED [DLL (=(())]]=(((ss))
(18)
which is equivalent minimizing loss function over Te:
((0) = न log no-
(19)
Our one-step policy target in Equation 6 approximates मे with 3 fixed to and Tprior regularized towards the behavior policy ns, yielding approximate regularized policy improvement.
14

---

Published as a conference paper at ICLR 2023
C MODEL-BASED OFFLINE REINFORCEMENT LEARNING
We discuss the related works in model-based offline reinforcement learning in this section. Model- based reinforcement learning refers to the class of methods that learn the dynamics function P( (5r) and optionally the reward function r(s, a, s'), which are usually utilized for planning. Levine al. (2020) has discussed several model-based offline RL algorithms in detail. The most
relevant works to ours include MOReL (Kidambi et al. 2020), MOPO (Yu et al. 2020 COMBO
Yu et al.
2021 and MuZero Unplugged (Schrittwieser et al. 2021).
MOReL (Kidambi et al. 2020) proposes to learn an ensemble of dynamics models from the offline dataset, and then utilize construct pessimistic-MDD (P-MDP), with which a normal RL agent can interact to collect experiences for learning. The construction of the P-MDP is based on the unknown state-action detector (USAD), which is realized by computing the ensemble discrepancy. If the discrepancy is larger than a threshold, then this state is treated as the absorbing state and taking the action will be given negative reward as the penalty. This would help constrain the policy not entering the unsafe regions that is not covered by the dataset.
MOPO (Yu al., 2020) takes a similar approach with MOReL by using uncertainty quantification to construct a lower bound for policy performance. The main difference is that a soft reward penalty constructed by an estimate of the model's error and the policy is then trained in the resulting MDP. COMBO YYu et al., (021)) further extends MOPO by avoiding explicit uncertainty quantification for incorporating conservatism. Instead, a critic function is learned using both the offline dataset and the synthetic model rollouts, where the conservatism is achieved by extending CQL to penalize the value function in model simulated state-action tuples that are not in the support of the offline dataset.
MOReL, MOPO and COMBO have theoretically shown that the policy's 's performance under the learned model bounds its performance in the real MDP, and achieved promising empirical results on state-based benchmarks such as D4RL (Fu et al,, 2020). However, it is unclear how such algo- rithmic frameworks can be transferred to image- based domains such as Atari in the RL Unplugged benchmark (Gulcehre et al. 2020)). Unlike state-based environments, learning the dynamics model for image-based tasks is challenging, and using the ensemble of learned dynamics help policy learning is even compute expensive, leading such algorithms not suitable for complex tasks.
Unlike methods discussed above, the Latent Dynamics Model (Section 2.2) does not aim to learn the environment dynamics explicitly and hence it does not require the reconstruction of the next state's observation, making more practical for image-based tasks. Furthermore, instead of the two-stage process of learning the MDP and planning with the learned MDP, the latent dynamics model facilitates the end-to-end training for learning the model and using the model. Both MuZero Unplugged (Schrittwieser et al. (02T) and ROSMO (ours) lie in this family of algorithms. MuZero Unplugged relies on the Monte-Card Tree Search to plan with the learned model for proposing learning targets, which have scrutinized in this paper and shown several deficiencies of (Sec-) [ion 3.1 Section 4.1). Motivated by the desiderata of model-based offline RL, including computa- tional efficiency, robustness compounding extrapolation errors and policy constraints, we devel- oped ROSMO, which uses a one-step look-ahead for policy improvement and incorporates behavior regularization for policy constraint. Our simpler algorithm is computationally efficient, and able to outperform MuZero Unplugged as well as other offline RL methods (including model-free and model-based" 41 the standard offline Atari benchmark Seccion 4.2).
D EXPERIMENTAL DETAILS
D.1 HARDWARE AND SOFTWARE
We use TPUv3-8 machines for all the experiments in Atari and use CPU servers with 60 cores for BSuite experiments. Our code is implemented using JAX (Bradbury et al. 2018).
stagestage model-based methods are compared in the Appendix F.1
15

---

Published as a conference paper at ICLR 2023
f-r
11
02
0
1+
4T
81
$1
$
(3) (b)
80
fo
mf.of
Q
hg
Ap
Pact <hy
Pre
Pish
(c)
(d)
Figure 5: (Left) Illustration on the training procedure: root node unrolling, recurrently applying dynamics func- tion, the model predictions and the improvement targets. (Right) Network architecture. (a) is for ROSMO and MuZero Unplugged; (b) is for Behavior Cloning; (c) is for Critic Regularized Regression; (d) is for Conserva- tive Q-Learning.
D.2 IMPLEMENTATION
D.2.1 NETWORK ARCHITECTURE
For both BSuite and Atari experiments, use network architecture based on the one used by MuZero Unplugged. In visual domains such as Atari games, the ResNet v2 style pre-activation residual blocks with layer normalization used model the representation, dynamics and pre- diction functions, while fully connected layers are used for simpler environments such as BSuite tasks.
We explain the network architecture for Atari in details. Figure 5(right) illustrates the network architectures used for the implementation of different algorithms. The blocks in light blue the representation function (in the terminology of ROSMO or MuZero Unplugged), which encodes the observation into a hidden state. The overall network size is downscaled compared to Schrittwieser et al. 2021) due to the experimentation cost. For the stacked grayscale image input of size 84 x 84 4, we firstly downsample as follows (with kernel size 3 for all convolutions):
convolution with stride 2 and 32 output channels.
residual block with 32 channels.
convolution with stride 2 and 64 output channels.
2 residual block with 64 channels.
Average pooling with stride 2.
residual block with channels.
Average pooling with stride 2.
Then 6 residual blocks are used to complete the representation function. We use 2 residual blocks for the dynamics function (blocks in yellow) as well as the prediction function (blocks in pink). All residual blocks are with 64 hidden channels. All the network blocks kept the same across different algorithms to ensure similar neural network capacity for a fair comparison.
For ROSMO and MuZero Unplugged, the input of the dynamics function is the latent state tiled with the one-hot encoded action vector. Two fully connected layers with 128 hidden units are used for the reward, value and policy predictions. The output size for policy the size of action space, while the output size for reward and value is the number of bins (601) used for the categorical representation described in Appendix A.2
For Behavior Cloning, we directly concatenate the representation function with the prediction func- tion model the policy network (Figure b)).ght-.)). For Critic Regularized Regression, the net- work is similar to our dynamics function but without the next state prediction, and the policy network based on our prediction function (Figure 5/right-c)). We follow the official code for the choice of hyperparameters, and we also employ the same categorical representation for Q value learning.
https:: thub. deepmii nd/ acme
16

---

Published as a conference paper at ICLR 2023
Parameter
Value
Frames stacked
4
Sticky action
True
Discount factor
0.9974
Batch size
512
Optimizer
Adamw
Optimizer learning rate
7 10
Optimizer weight decay
10
Learning rate decay rate
0.1
Max gradient norm
5
Target network update interval
200
Policy loss coefficient
1
Value loss coefficient
0.25
Unroll length
5
TD steps
5
Bin size
601
Table Atari hyperparameters shared by ROSMO and MuZero Unplugged.
Parameter
Value
Representation MIP
[64, 64, 32]
Dynamics MLP
[32, 256, 32]
Prediction MLP
[32]
Discount factor
0.9974
Batch size
128
Optimizer
Adamw
Optimizer learning rate
10-1
Optimizer weight decay Learning rate decay rate Max gradient norm Target network update interval Policy loss coefficient
10-4
0.1
5
200
1
Value loss coefficient
0.25
Unroll length
5
TD steps Bin size
3
20
Table 4: BSuite hyperparameters shared by ROSMO and MuZero Unplugged.
For Conservative Q-Learning, similarly, we follow the official code and their hyperparameters, but with our network illustrated in Figure 5(right-d).
D.2.2 EXPERIMENT SETTINGS
The hyperparameters shared by ROSMO and MuZero Unplugged for Atari environments is given in Table 3] and that for BSuite environments is given in Table 4| In addition, the behavior regularization strength (a) used in ROSMO is chosen to be 0.2. The simulation budget for MuZero Unplugged is 20 for Atari and for BSuite, and the depth is not limited. We use the official library to implement the MCTS used in MuZero Unplugged and its parameters follow the original settings (Schrittwieser et al. 2020).
For all experiments in Atari we use 3 seeds, and comparison between ROSMO and MuZero Unplugged is made apple-to-apple by only replacing the policy and value targets as well as the regularizer.
seeds for all BSuite experiments. The
https:://qithub.ccmmm// ral 10mna22907 https://gitthbb.com nd/ mctx
17

---

Published as a conference paper at ICLR 2023
E BSUITE DATASET
We follow the setup of Gulcehre al. (2021) to generate episodic trajectory data by training DQN agents for three tasks: cartpole, catch and mountain_car We also add stochastic noise to the origi- nally deterministic environments by randomly replacing the agent action with a uniformly sampled action with probability of E {0, 0.1, 0.3, 0.5}. We use env envlogge trajectories through the training process. More details of the episodic dataset are provided in Table We also record the score of random policy and an online DQN agent the three environments in which can be used to normalize the episode return for evaluation.
record complete episode
Table 6
Environments
Number of episodes
Number of transitions
Average episode return
cartpole (e 0.0)
1,000
630,262
629.71
cartpole 0.1)
1,000
779,491
779.01
cartpole (c 0.3)
1,000
787,350
786.86
cartpole (E 0.5)
1,000
527,528
526.75
catch (c = 0.0)
2,000
18,000
0.71
catch (e 0.1)
2,000
18,000
0,60
catch (c = 0.3)
2,000
18,000
0.25
catch 0.5)
2,000
18,000
-0.04
mountain car 0.0)
500
82,342
-164.68
mountain.car (c 0.1)
500
147,116
-24.23
mountain car 0.3)
500
138,262
-276.52
mountain car (e 0..5)
500
167,688 06,688
-335.37
Table 5: BSuite episodic dataset details.
random agent
online DQN agent
cartpole catch
64.83
1,001.00
1.00 -0.66
mountain ccar
-1,000.00
-102.16
Table Episode return random and online agent on BSuite environments.
F ADDITIONAL EMPIRICAL RESULTS
F.1 BENCHMARK RESULTS FOR MODEL-BASED METHODS
In Section 4.2
we have compared ROSMO with Behavior Cloning, Conservative Q-Learning, Critic Regularized Regression and MuZero Unplugged on the offline Atari benchmark and presented the results in Figure 3) The comparison was made apple-to tople-to-apple as we standardized the neural network architecture and the optimization steps (Appendix D.2). However, comparing other model-based offline methods such as MOReL (Kidambi et al. (020) and COMBO (Yu et al., (202)) in Figure 3 less feasible. This is because these methods adopt two-stage training procedure, where they first learn the MDP and then plan with the learned MDP. Moreover, MOReL and COMBO are not readily suitable for the Atari benchmark we use (MOReL only focuses on state-based tasks and COMBO's implementation and dataset are not released). Therefore, we implemented MORel. and COMBO and compared them with the other methods. For both of them, we have tuned the hyper- parameters to report the results of the best configuratio.. In the following sections we introduce our implementation details and report our experimental results.
F.1.1 DYNAMICS MODEL
Both MOReL and COMBO need to learn the dynamics model in the first stage. Since work- ing with image-based tasks, we use the Dreamer VreyyerV2-stylee (Hafner et al.) 2021) framework to learn
httpps:/ http://ggtttt.u.tto thub. cwm deepmi nd/envllggee ogger
18

---

Published as a conference paper at ICLR 2023
the dynamics model. In particular, only enable the world model learning in the DreamerV2, com- posed of the Recurrent Spaceeppaeeeee Model, and the image, reward and discount predictor. We adapt the pydreame of code base and follow the default sypee-parameters to train until convergence. We also trained an ensemble of dynamics models for uncertainty quantification.
F.1.2 MOREL
Given the trained ensemble of dynamics models f2..... }. MOReL constructs a pessimistic MDP (P-MDP) by cross validating the collected trajectories in each learned model and comput- ing the ensemble discrepancy as disc(s, a) = max.,; [[f a) - (s, a) disc(s, a) is larger than certain threshold, the state is regarded the terminal state and r(s, a) is assigned small value as penalty. Since the MOReL. framework is agnostic to the planner, we employ a PPO (Schulman al. 2017) agent adapted from open-source implementation" with strong online RI. performance to learn from the P-MDP.
F.1.3 COMBO
We refer to the implementation of a public repository OfflineRI to adapt COMBO for Atari Games. We use the trained DreamerV2 generate rollout data, where the dynamics model is randomly chosen from the ensemble for each rollout procedure. The rollout length is set to 10. The policy training part very similar CQL, where the only difference lies in that COMBO takes mix of offline dataset and rollout dataset as input. The ratio between model rollouts and offline data set to 0.5. The other hyperparameters just follow our CQL implementation.
F.1.4 RESULTS
Following the setting in ablation study, choose MsPacman representative game train both MOReL and COMBO with shared world models. We run the experiments with different random seeds and report the mean and standard deviation. Table 7 shows that among all the model- based methods, ROSMO achieves the best result on MsPacman.
MOReL
COMBO
MZU
ROSMO
MsPacman
1476.476 476.4478.mmm sarar
1538.33
4539.045 TOU
5019.762
Table 7: Episode return on MsPacman of different model-hased offline RI algorithms.
19

---

Published as a conference paper at ICLR 2023
F.2 LEARNING CURVES OF INDIVIDUAL ATARI GAMES
Figure 6 shows the learning curves in terms of IQM episode return for individual Atari games to compare ROSMO with other baseline methods, as complementary results for Figure Table
records the numerical results of the IQM episode return of individual Atari games.
Amidar
Gravitar
Pung
20
IOM ephode retum
10-
1.00
800
400
0
-20-
Astortx
Jamesbond
Qbort
I0M spuode return
800
20000
600
400
10000
5000
Breakout
10M epeodde returrs
400
15000
300
4000
10000
2000
1000
5000
0
Prostbite
Phoenk
Seaquest
I2M epocdd retur.
6000
3000
20000
4000
1000
10000
2000
50
100
150
200
0
200
150
200
Number of Stops (In thousands)
Number of Staps (in thousands)
Number of Steas On thousands)
R0SMO
HZU
ORR
COL
BC
Figure 6: IQM episode return of different algorithms on individual Atari games.
BC
CRR
CQL
MZU
ROSMO
Amidar
30.381
22.143ana
51.286
76.452
38.405
Asterix
2633.333.swa
6344.048
26890.476
476
29061.905.waa
16740.476maaan
Breakout
85.143mm
218.952 11577
390.119
440.905
Frostbite
586.16
686.19mmar
2854.286-maa
3337.38
4051.19.
3996. 19 19
Gravitar
400.0.w.ss
345.238
52.381
792.857 7
753.571 -orson
Jamesbond
402.38
Lawar
513.095
102.38
602.38 1 LUEUT
639. $29.2866.um Laress
MsPacman
2733.30 333
3736.905
2141.667
$539.048.mmm
5019. 762.762... LEaTTS
Phoenix
5901. .9.5..mm.
5954.286 a2
3510.238
6103.81 1.:raa..
21550.476 6ussesa
Pong
14.976
19.095
18.762
18.452
20.452
Qbert
7497.6 11997.69... A5221422
12618.457 4522. 226.770
13114,22 24. Lerraa
13121.422
15848,81 81, 1000
Riverraid
961 4..18.444........ -rasnn
12279.
762 202
14887. 1 143 43-ma
15156.905..m..
19399.28 286
Seaquest
1087. 143
3879.048....
wa.
7237.619.man.
6775.238.mm
6642.857
Table 8: Numerical results of the IQM episode return of individual Atari games.
20

---

Published as a conference paper at ICLR 2023
F.3 COMPARISON ON ROSMO AND ONESTEP
Figure 7 shows the training curves of ROSMO and the OneStep variant (removing the behaviour regularization term defined in Equation 11). These results extend the ablation in Figure 4jd) with longer training time and more games. The comparison shows that ROSMO is able to achieve faster convergence, while resulting in similar slightly better final performance. We further conducted experiments with only 1% data to verify the effect of behavior regularization when the data coverage low. As shown in Figure 8 with limited data, ROSMO performs significantly better compared to OneStep. We can also observe that both methods suffer from over-fitting when the training goes longer due to limited sample size. To effectively handle this issue, we could resort policy evalua- tion to early stop the training and select the best trained policy. Ideally in offline RL we need to use offline policy evaluation methods (Voloshin al. 2021) for this purpose, which is unfortunately not trivial for difficult tasks. We leave this for future research since it is beyond the scope of this paper.
Asterix
MsPacman
Qbort
IQM eplode return
30000
6000
15000
5000
200000
4000
10000
3000
10000
2000
5000
1000
0
0
Pong
Seaquest
20
8000
5000
10M egisode reture
400
10
6000
300
4000
200
2000
100
-20
150
200
100
100
150
200
Number of Steps (in thousands)
Number Steps (in thousands)
Number of Steps On thousands)
ROSMO
OneStep
Figure 7: IQM episode return ROSMO and OneStep.
Asterix
MsPacman
Obert
10000
2500
3000
(QM episode return
8000
2000
2000
1500
6000
1000
1000
2000
500
0
Breakout
Pong
Seaquest
10M episode retur
250
3000
200
150
0-
2000
100
-10
1000
0
0
20
40
60
80
20
40
60
80
20
40
60
80
Number of Steps (in thousands)
Number of Steps (in thousands)
Number of Steps (in thousands)
ROSMO
OneStep
Figure& 8: 1QM episode return of ROSMO and OneStep trained only with 1 %% data.
21

---

Published as a conference paper at ICLR 2023
F.4 LEARNING CURVES FOR BSUITE EXPERIMENTS
Figure 9 and Figure 10 show the learning curves of individual settings in our BSuite analysis for noisy data model capacity.
1.0
0.5
0.0
-0.5
ROSMO
MZU
Figure 9: Learning curves of ROSMO and MuZero Unplugged on noisy catch environment complementary results for Table
240
-430
-3095
253
282
100
Figure 10: Learning curves of ROSMO and MuZero Unplugged on individual BSuite environment at different dynamics capacity, complementary results for Figure 2 b).
22


Reviews:
[{"review_id": "j5hPltrY6Rp", "reviewer": "Reviewer_tLyk", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper studies offline model-based RL, specifically identified several shortcomings of the previously proposed MuZero Unplugged. The proposed ROSMO (Regularized One-Step Model-based Offline RL) incorporates (i) one-step lookahead instead of MCTS planning, (ii) behavior regularization. \n\nExtensive experiments and ablation studies are conducted on the Atari benchmark and achieve impressive results, both in terms of performance and runtime efficiency. ", "summary_of_the_review": "Overall the paper is written quite clearly and empirical results are indeed strong, however my slight concern is where this contribution lies within the literature, perhaps a dedicated related works section would help - which should include more discussion on previous model-based offline RL approaches in \"Offline RL Tutorial by Levine et al. 2020\". ", "strength_and_weaknesses": "**Strengths**\n- Empirical results are pretty extensive and strong across the board. \n- Sec 4.1 ablations help the reader understand what each modeling component addresses. \n\n**Weaknesses**\n- MOPO/MOReL: \"both of them focus on state-based control tasks and hence are not compared here\" - could you clarify what this mean? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "OGZyzzFSsVQ", "reviewer": "Reviewer_Rze2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper introduces ROSMO, a simpler and better version of Muzero Unplugged for offline policy optimization. In the offline setting, Muzero Unplugged can suffer from limited data-coverage, inaccurate models and an expensive compute budget. Instead of relying on an expensive MCTS and suffering from out of distribution errors as well as compounding of errors, ROSMO resorts to a one step lookahead approach to learn the same components as the one needed by Muzero Unplugged (policy and value targets) and a behavior cloning regularization to prevent going into an unsafe regions. ROSMO outperforms Muzero Unplugged on the Bsuite and Atari benchmarks. The authors also provide an ablation study.", "summary_of_the_review": "Overall I think this is a nice paper and good contribution to the community. A comparison to other model-based algorithms for offline RL on the D4RL would have been a big plus (Morel/MOPO, ...).", "strength_and_weaknesses": "The paper is well written and well organized. It was easy to read and understand. I liked that the authors recalls the ideas behind Muzero Unplugged in section 2.2 (in a more general way) and the MCTS policy improvement used by Muzero Unplugged in section 2.3, I think this eases the reading of the paper for people not very familiar with Muzero Unplugged.\n\nMuzero Unplugged is a not so simple algorithm to implement and has no official open sourced implementation as far as I know. I agree with the authors write that \"Muzero Unplugged is sophisticated\". Having a simpler, yet as effective (and even better) algorithm is a great contribution for the community, especially if the implementation is to be made open sources after acceptance of the paper.\n\nEmpirical evaluations of Muzero Unplugged and ROSMO as well as the ablation study show the pertinence and the superiority of the ROSMO algorithm.\n\nIt would be better to say that having a one-step update might not be the only way to solve the compounding of errors as well as the fact that Muzero can go in unsafe regions. For instance having a model with its own uncertainty could be another direction to avoid unsafe regions and planning too much into the future. I think that planning (with MCTS or a different search algorithm) could make better targets if the uncertainty is used correctly. I agree that this would make a more complicated algorithm and for that it's good to have the ROSMO algorithm as it is simpler and will provide a good baseline when developing other algorithms.\n\nThe part on the stochastic transitions in section 4.1 is the least clear for me: why not add noise in the observations? why is it connected to the compounding of errors?\n\nDesigning simpler algorithms is a nice contribution to the community and for this reason I would be in favor of accepting the paper. \n\nIt is a pity that the authors are not using the D4RL benchmark which has become the default benchmark for offline RL and for which performances of MOPO/CQP/Morel etc... are available. It would have been nice to compare the ROSMO algorithm to these algorithms. Muzero unplugged is said to be the state of the art but was not compared to these algorithms to the best of my knowledge which are uncertainty based.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "1Ktn0PWGmi", "reviewer": "Reviewer_Dxeb", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper presents ROSMO, a model-based offline RL algorithm based on MuZero Unplugged (MZU). ROSMO improves MZU in the offline RL setting in two main ways. First, it only performs one-step rollout for advantage estimation, rather than MCTS with multi-step rollouts. Second, it uses advantage-based policy improvement with advantage-filtered behavior cloning. ROSMO outperforms MZU in a low-data regime and stochastic transition and is less sensitive to network parameters. In the Atari benchmark, ROSMO outperforms offline RL baseline algorithms.\n", "summary_of_the_review": "The paper provides modifications to MuZero Unplugged for offline RL, which makes the algorithm simpler yet more effective and this is a nice contribution. Still, the technical novelty is a bit limited, and it would be great to have more offline RL baselines in the experiments.", "strength_and_weaknesses": "[Strengths]\n1. The paper is well-written and easy to follow. The modification to MZU for offline RL is well-motivated, and it makes the algorithm simpler yet more effective in terms of computation cost and performance.\n2. ROSMO outperforms baseline algorithms in the experiments.\n\n\n[Weaknesses]\n1. Technical novelty is a bit limited. The short-horizon rollout for MBRL has already been widely used (e.g., [1]), and the behavior regularization is also very similar to the one in previous works (e.g. [2,3]). I am curious to see the result when only (11) is used for the policy improvement while not using (9), since it seems (11) can also serve as a policy improvement.\n2. There are missing comparison with model-based offline RL algorithms in the experiments, e.g. [4,5,6].\n3. In (1), there is no explicit definition of each term, i.e. $\\ell^r$, $\\ell^v$, and $\\ell^p$.\n4. In Figure 4d, it seems OneStep is not yet converged. It would be great to provide the result with a longer number of steps to see whether ROSMO clearly outperforms OneStep at convergence. In the current result of Figure 4d, the conclusion is unclear to me.\n5. It would be great to include the result of baseline algorithms in Figure 2. Why is ROSMO only being compared with MZU?\n6. In section 4.1.(3), why is MZU more sensitive to the parametrization of dynamics? It would be great to provide some possible reasons.\n\n\n[1] Janner et al., When to Trust Your Model: Model-Based Policy Optimization, 2019\n\n[2] Oh et al., Self-Imitation Learning, 2018\n\n[3] Wang et al., Critic Regularized Regression, 2020\n\n[4] Yu et al., COMBO: Conservative Offline Model-Based Policy Optimization, 2021\n\n[5] Kidambi et al., MOReL : Model-Based Offline Reinforcement Learning, 2020\n\n[6] Yu et al., MOPO: Model-based Offline Policy Optimization, 2020\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster