# Model-Augmented $Q$-Learning

Youngmin Oh
Samsung Advanced Institute of Technology  Jinwoo Shin Eunho Yang
Korea Advanced Institute of Science and Technology  Sung Ju Hwang
Korea Advanced Institute of Science and Technology
<EMAIL_ADDRESS>
###### Abstract

In recent years, $Q$-learning has become indispensable for model-free
reinforcement learning (MFRL). However, it suffers from well-known problems
such as under- and overestimation bias of the value, which may adversely
affect the policy learning. To resolve this issue, we propose a MFRL framework
that is augmented with the components of model-based RL. Specifically, we
propose to estimate not only the $Q$-values but also both the transition and
the reward with a shared network. We further utilize the estimated reward from
the model estimators for $Q$-learning, which promotes interaction between the
estimators. We show that the proposed scheme, called Model-augmented
$Q$-learning (MQL), obtains a policy-invariant solution which is identical to
the solution obtained by learning with true reward. Finally, we also provide a
trick to prioritize past experiences in the replay buffer by utilizing model-
estimation errors. We experimentally validate MQL built upon state-of-the-art
off-policy MFRL methods, and show that MQL largely improves their performance
and convergence. The proposed scheme is simple to implement and does not
require additional training cost.

##  1 Introduction

Model-free reinforcement learning (MFRL), which does not learn a model about
the environment, has achieved remarkable success on various tasks, e.g., Atari
games (Schaul et al.,, 2016; Mnih et al.,, 2015; Hessel et al.,, 2018; van
Hasselt et al.,, 2019), Starcraft (Vinyals et al.,, 2019), and robotic control
tasks (Haarnoja et al.,, 2018b; Fujimoto et al.,, 2018; Schulman et al.,,
2017, 2015), due to its simplicity that eliminates the needs of hand-designing
or learning a model, and many techniques that enable sample-efficient learning
without a model.

Since MFRL algorithms learn agents’ policy without environment models,
$Q$-learning is critical for policy learning. Some recent MFRL methods such as
Rainbow (Hessel et al.,, 2018; van Hasselt et al.,, 2019) achieve impressive
performance with $Q$-learning without learning the policy networks. MFRL
algorithms (Haarnoja et al.,, 2018b, a; Fujimoto et al.,, 2018) with actor-
critic models (Konda & Tsitsiklis,, 2000), which learns the policy with the
actor, still relies on the critic that estimates the $Q$-values to guide its
learning. Here, an accurate estimation of the values of the state-action pairs
with $Q$-learning is an essential part of many MFRL algorithms for their
success, since if they are over- or underestimated, then the bias will
adversely affect the policy learning (Hasselt,, 2010; Fujimoto et al.,, 2018;
Lan et al.,, 2019).

However, the lack of model in MFRL makes them sample-inefficient compared to
model-based RL, which can generate samples from the model of the environment
that is either manually defined or learned. Many prior works (Janner et al.,,
2019; Nagabandi et al.,, 2018) have proposed to combine the two methods in a
single framework, which first train a model-based controller and collects
trajectories to initializes policy network that continually learns using the
trajectories from the controller.111We provide more detailed literature survey
in the supplementary material.

Yet, instead of utilizing the generated transitions as in those methods, we
propose to use the components of the model-based RL to _augment_ the learning
of model-free RL. Specifically, we note that the reward $\mathcal{R}(s,a)$,
and the transition $\mathcal{T}(s,a)$ are functions that have the same input
space as the $Q$-function $Q(s,a)$, and modify the $Q$-network into a _Model-
augmented $Q$-network_ (MQN) which additionally predicts the reward and the
transition function (see Figure 1). However, instead of naively performing
multi-task learning via shared network weights, we promote interactions
between the $Q$-learner and the model estimators. Specifically, we use the
estimated reward from the model (MReward) to calculate the temporal difference
(TD) target for $Q$-learning, and use the estimated $Q$-value to learn the
reward. Moreover, we further utilize the estimation errors for the reward and
transition function to prioritize the past experiences with high errors from
the replay buffer, which we say _Model-augmented Prioritized Experience
Replay_ (MPER).

---
Figure 1: Overview of Model-Augmented $Q$-Learning. Our proposed method
improves upon the MFRL algorithms. First, we modify a given $Q$-network to a
model-augmented $Q$-network (MQN), which also learns the environment model via
parameter sharing. The estimators interact with each other via Model-augmented
TD learning, where the TD errors of MQN are calculated using the estimated
reward (MReward) as well as the transition map. Further, we use the model
estimation errors compute priorities of experiences for sampling transitions
in a buffer.

Model-augmented $Q$-learning has the following advantages over the existing
MFRL algorithms on various environments. First, learning to predict $Q$-values
in the long-term is difficult in general, and thus estimating the short-term
reward with the shared model helps it learn both short-term and long-term
viewpoints. Second, the estimators are highly related in the way that they
help with the others’ prediction task. Specifically, by utilizing the
estimated reward and the transition, we can calculate the $Q$-values for a
given policy _without_ interacting with the environment. Third, the proposed
MPER can sample more diverse experiences to improve the sample-efficiency for
not only MQN but also policy networks since we set the sampling priority as a
weighted sum of the TD and the model estimation errors.

In summary, our contributions are followings:

  * •

We propose _Model-Augmented Q-Learning_ (MQL), which not only augments the
conventional $Q$-learning with the model estimation tasks via parameter
sharing, but also utilizes the estimated values in each task to the
computation of the objective in the other task.

  * •

We show that MQL satisfies policy invariance and converges to an identical
solution obtainable with $Q$-learning with the true reward.

  * •

We further propose to use the model estimation errors when prioritizing the
the past transitions to sample in the experience replay (MPER).

  * •

We show that MQL is applicable to any modern MFRL algorithms (e.g., SAC, TD3,
and Rainbow) with neural network estimators, and significantly improves their
performance and sample-efficiency under various environments: MQL is easy-to-
implement and does not incur additional training cost.

##  2 Method

In this section, we first introduce model-free reinforcement learning (MFRL)
with the actor-critic architecture and prioritized experience replay (PER).
Then we propose our model-augmented $Q$-network (MQN) and define model-
augmented $Q$-learning (MQL) by proposing the model-augmented reward (MReward)
and the corresponding TD learning objective with the MReward. By utilizing
MQN’s TD-error equation (Eq. (13)), we promote interactions between MQN’s
estimators, such that estimating the models helps better estimate the
$Q$-value, and vice versa. We then prove that our method satisfies policy
invariance. Finally, we present the model-augmented prioritized experience
replay (MPER) which samples the past experiences based on the model estimation
errors.

###  2.1 Preliminary

Model-Free Reinforcement Learning (MFRL). In a standard reinforcement learning
framework, an agent interacts with a given environment over discrete timesteps
$t$. The agent selects action $a_{t}\in A$ by its policy $\pi$ on the current
state $s_{t}\in S$ to receive the next state $s_{t+1}\in S$ and a reward
$r_{t}$ from the environment, where $A$ and $B$ are action and state spaces,
respectively. The agent’s objective is to learn a policy $\pi$ that maximizes
the return $R_{t}=\sum_{k=0}^{\infty}\gamma^{k}r_{t+k}$, which is the
discounted cumulative rewards from the timestep $t$ with a discount factor of
$\gamma\in[0,1)$, at each state $s_{t}$.

Although our method is applicable to $Q$-learning in general, throughout this
section, we focus on off-policy actor-critic RL algorithms with an experience
replay buffer $\mathcal{B}$, which consist of the policy (i.e., actor)
$\pi_{\Theta}(a|s)$ and $Q$-function networks (i.e., critic)
$Q_{\theta}(s,a)$, where $\Theta$ and $\theta$ are their parameters,
respectively. The most commonly used loss for $Q_{\theta}(s,a)$ (Mnih et al.,,
2015; Haarnoja et al.,, 2018b; Fujimoto et al.,, 2018) is

$\displaystyle\mathbb{E}_{(s_{t},a_{t},r_{t},s_{t+1})\sim\mathcal{B}}\left[\|\delta_{t}Q_{\theta}^{\pi_{\Theta}}\|_{2}^{2}\right],$
(1)

where $\|\cdot\|_{2}^{2}$ is the mean-square error and $\delta_{t}Q_{\theta}$
is the Temporal Difference error (TD-error) defined as follows: for a given
transition $(s_{t},a_{t},r_{t},s_{t+1})$, its TD-error $\delta_{t}Q_{\theta}$
is:

$\displaystyle\delta_{t}Q_{\theta}^{\pi_{\Theta}}=\delta
Q_{\theta}^{\pi_{\Theta}}(s_{t},a_{t},r_{t},s_{t+1})$
$\displaystyle=r_{t}+\gamma
Q_{\theta}\left(s_{t+1},a^{\prime}\right)-Q_{\theta}\left(s_{t},a_{t}\right),$
(2)

where $a^{\prime}\sim\pi_{\Theta}(\cdot|s_{t+1})$. We can interpret the TD-
error as a measure of how surprising or ‘unexpected’ the transition is.

Prioritized Experience Replay. We now explain the concept of PER since we will
later augment it with model-based components. Let $[n]$ be defined as set
$\\{1,\cdots,n\\}$ for a positive integer $n$. Without loss of generality, we
can suppose that the replay buffer $\mathcal{B}$ stores the following
information as its $i$-th transition:

$\displaystyle\mathcal{B}_{i}=\left(s_{\kappa(i)},a_{\kappa(i)},r_{\kappa(i)},s_{\kappa(i)+1}\right),$
(3)

with a function from the index of $\mathcal{B}$ to a corresponding timestep:
$\kappa\left(i\right)$. PER calculates each $\mathcal{B}_{i}$’s priority
$\sigma_{i}$ as the recently computed TD-error about itself and collects a set
of priority scores:

$\displaystyle\mathcal{P}_{\mathcal{B}}=\\{\sigma_{1},\cdots,\sigma_{\left|\mathcal{B}\right|}\\},$
(4)

where each priority $\sigma_{i}$ is updated whenever the corresponding
transition is sampled for training the actor and critic networks. The TD-error
(2) is the most frequently used quantity to make the priority set (4) (Schaul
et al.,, 2016; Brittain et al.,, 2019; Hessel et al.,, 2018; van Hasselt et
al.,, 2019). Then the sampling strategy of PER is to determine an index set
$I$ in $\left[\left|\mathcal{B}\right|\right]$ from the probability $p_{i}$ of
$i$-th transition defined by the priority set:

$p_{i}=\frac{\sigma^{\alpha}_{i}}{\sum_{k\in\left[\left|\mathcal{B}\right|\right]}\sigma^{\alpha}_{k}},$
(5)

with a hyper-parameter $\alpha>0$. Next, to compensate the bias of
probabilities, we calculate the importance-sampling weights as follows:

$w_{i}=\left(\frac{1}{|\mathcal{B}|p(i)}\right)^{\beta},$ (6)

where $\beta>0$ is also a hyper-parameter.

Algorithm 1 MQN, MPER, MReward to actor-critic RL methods

Initialize the network $\mathcal{M}$ parameters $\theta$, the actor parameters
$\Theta$, a replay buffer $\mathcal{B}\leftarrow\emptyset$, priority set
$\mathcal{P}_{\mathcal{B}}\leftarrow\emptyset$, and index set
$\mathcal{I}\leftarrow\emptyset$.

Set the batch size $m$.

for each timestep $t$ do

Choose $a_{t}$ from the actor and collect a transition
$(s_{t},a_{t},r_{t},s_{t+1})$ from the environment.

Update replay buffer
$\mathcal{B}\leftarrow\mathcal{B}\cup\\{(s_{t},a_{t},r_{t},s_{t+1})\\}$ and
priority set
$\mathcal{P}_{\mathcal{B}}\leftarrow\mathcal{P}_{\mathcal{B}}\cup\\{1.0\\}$.

for each gradient step do

Sample an index $I$ by the given set $\mathcal{P}_{\mathcal{B}}$ and Eq. (5)
with $|I|=m$.

Calculate weights $\\{w_{i}\\}_{i\in I}$ and a priority set
$\\{\sigma_{k}\\}_{k\in I}$ by (6) and (15), respectively.

Train
$\mathcal{M}_{\theta}=(\widetilde{Q}_{\theta},\mathcal{R}_{\theta},\mathcal{T}_{\theta})$
by (14) and the actor by $\widetilde{Q}_{\theta}$ with batch
$\\{\mathcal{B}_{i}\\}_{i\in I}\subset\mathcal{B}$ and corresponding weights
$\\{w_{i}\\}_{i\in I}$.

end for

end for

###  2.2 Model-augmented $Q$-learning

Model-augmented $Q$-network. In an RL framework, an environment model consists
of two maps: a reward map $\mathcal{R}(s,a)$ and a transition model
$\mathcal{T}(s,a).$ We observe that these two maps and $Q_{\theta}(s,a)$ have
the same input domain, i.e., $S\times A$, where $S$ and $A$ are state and
action spaces. Motivated by this, we slightly modify $Q_{\theta}(s,a)$ as
$\mathcal{M}_{\theta}(s,a)$ to additionally predict the environment model: a
reward network $\mathcal{R}_{\theta}(s,a)$ and a transition model
$\mathcal{T}_{\theta}(s,a)$ in parallel via parameter sharing. We refer to
$\mathcal{M}_{\theta}(s,a)$ as a _Model-augmented $Q$-network_ (MQN) such that

$\mathcal{M}_{\theta}=\left(\widetilde{Q}_{\theta},\mathcal{R}_{\theta},\mathcal{T}_{\theta}\right),$

where $\widetilde{Q}_{\theta}$ measures $Q$-values with MReward with the
additional estimators, which we will explain soon below.

Model-Augmented $Q$-Learning (MQL). Instead of simply resorting to weight
sharing, we formulate a new objective for MQN to promote the interaction
between the estimators. To this end, we first propose the following losses the
estimation of $\mathcal{R}_{\theta}$ and $\mathcal{T}_{\theta}$:

$\displaystyle L_{\mathcal{R}_{\theta}}$
$\displaystyle=\mathbb{E}_{(s_{t},a_{t},r_{t},s_{t+1})\sim\mathcal{B}}\left[\|\delta_{t}\mathcal{R}_{\theta}\|_{2}^{2}\right],$
(7) $\displaystyle L_{\mathcal{T}{\theta}}$
$\displaystyle=\mathbb{E}_{(s_{t},a_{t},r_{t},s_{t+1})\sim\mathcal{B}}\left[\|\delta_{t}\mathcal{T}_{\theta}\|_{2}^{2}\right],$
(8)

where

$\displaystyle\delta_{t}\mathcal{R}_{\theta}$
$\displaystyle=\mathcal{R}_{\theta}(s_{t},a_{t})-r_{t},$ (9)
$\displaystyle\delta_{t}\mathcal{T}_{\theta}$
$\displaystyle=\mathcal{T}_{\theta}(s_{t},a_{t})-s_{t+1},$ (10)

Then by using the estimation errors in (9)-(10), we shape the new reward
$\widetilde{r}_{t}$ (MReward) instead of using $r_{t}$:

$\displaystyle\widetilde{r}_{t}$
$\displaystyle=\mathcal{R}_{\theta}+\zeta_{1}\delta_{t}\mathcal{R}_{\theta}+\zeta_{2}\delta_{t}\mathcal{T}_{\theta}.$
(11)

with some positive coefficients $\zeta_{1}$ and $\zeta_{2}.$ We use MReward to
develop a new TD-error about MQN, which encourage agent’s exploration compared
to the original $Q$-network. However, using MReward may change the direction
of agent’s learning. However, we will prove that learning of MQN with MReward
does not change the set of optimal policies in the last paragraph of this
section.

We formulate the loss for $\widetilde{Q}_{\theta}$ using the MReward (11),
that is similar to the TD-error in (2):

$\displaystyle
L_{\widetilde{Q}_{\theta}}=\mathbb{E}_{(s_{t},a_{t},s_{t+1})\sim\mathcal{B}}\left[\|\delta_{t}\widetilde{Q}_{\theta}^{\pi_{\Theta}}\|_{L^{2}}^{2}\right],$
(12)

where

$\displaystyle{\delta}_{t}\widetilde{Q}_{\theta}^{\pi_{\Theta}}={\delta}\widetilde{Q}_{\theta}^{\pi_{\Theta}}(s_{t},a_{t},r_{t},s_{t+1})$
$\displaystyle=\widetilde{Q}_{\theta}(s_{t},a_{t})-\left(\widetilde{r}_{t}+\gamma\widetilde{Q}_{\theta}(s_{t+1},a^{\prime})\right).$
(13)

Here $a^{\prime}$ denotes the action sampled by the current policy on the
current state, i.e., $a^{\prime}\sim\pi_{\Theta}(\cdot|s_{t+1})$. Then
combining (7)-(8) and (12), we obtain the loss for Model-augmented
$Q$-learning (MQL) of $\mathcal{M}_{\theta}$:

$\displaystyle
L_{\mathcal{M}_{\theta}}=\xi_{1}L_{\widetilde{Q}_{\theta}}+\xi_{2}L_{\mathcal{R}_{\theta}}+\xi_{3}L_{\mathcal{T}_{\theta}},$
(14)

with some positive coefficients $\xi_{1}$, $\xi_{2}$, and $\xi_{3}$. The
coefficients adaptively changes by the quantity of each loss by employing a
dynamic method in (Liang & Zhang,, 2020). We will specify how to determine the
coefficients in the supplementary material.

Model-augmented PER. Using (13), we make the estimators in MQN interact
positively with each other by proposing a suitable sampling method. we
formulate MPER from PER that compute the sampling probability of each
transition based on a single metric, e.g., TD-error. Since our method uses
$\mathcal{M}_{\theta}(s,a)$ instead of the original action-value function, we
modify the rule in obtaining the priority set in (4) for experiences in the
buffer $\mathcal{B}$ accordingly. To this end, we compute the priority of each
transition as the weighted sum of the TD-errors and model errors using (2) and
(9)-(10):

$\sigma_{i}=\xi_{1}|\delta_{\kappa(i)}\widetilde{Q}_{\theta}|^{2}+\xi_{2}|\delta_{\kappa(i)}\mathcal{R}_{\theta}|^{2}+\xi_{3}|\delta_{\kappa(i)}\mathcal{T}_{\theta}|^{2}.$
(15)

Then MPER computes probabilities and weights for corresponding transitions by
(5) and (6). When using the priority scoring in (15), the buffer can sample
various experiences that are useful for both long-term and short-term
viewpoints.

Interaction of MQN’s estimators. Now then we can state that the equation (13)
makes MQN’s estimators positively interact each other with MPER. For instance,
if the model errors are large, then $\delta_{t}Q_{\theta}^{\pi_{\Theta}}$ in
(13) should be also large since MReward $\widetilde{r}_{t}$ deviates from the
real reward. Since both errors are large, MPER samples experiences that are
valuable to learn both the estimators adaptively by the priority set. In other
words, due to the interaction between the two estimators in MQN, reducing
model errors also results in reducing TD errors simultaneously. The
interaction is not possible without both parameter sharing and the equation
(13), since it is nontrivial to select experiences that are valuable to all
estimators. We validate our statement in Section 3.1.

The detailed descriptions of our model-augmented $Q$-learning is provided in
Algorithm 1.

Convergence of MQL We establish the following guarantee on convergence of
Model-augmented $Q$-learning to a unique solution that is identical to the
solution obtained with the true reward, under a fixed policy:

######  Theorem 1.

Let $\mathcal{M}=\left(\widetilde{Q},\mathcal{R},\mathcal{T}\right)$ be
estimators of $Q$-value, reward, and transition maps defined as in Section
2.2. If $\mathcal{M}$ is updated by the following loss (14) with the MReward
(11) for a given policy $\pi^{*}$, then $\mathcal{M}$ converges to a unique
$\left(\widetilde{Q}^{*},\mathcal{R}^{*},\mathcal{T}^{*}\right)$ satisfying

$\displaystyle\widetilde{Q}^{*}(s_{t},a_{t})=r_{t}+\gamma\widetilde{Q}^{*}(s_{t+1},a)$
(16)

for any transition $(s_{t},a_{t},r_{t},s_{t+1})$, where
$a\sim\pi^{*}(\cdot|s_{t+1})$.

Theorem 1 suggests that the convergence of $\widetilde{Q}$ is invariant with
respect to MReward (11). In other words, our method guarantees policy
invariance, which means that the modification of the $Q$-network does not
affect a set of Pareto optimal policies. We provide the proof of this theorem
in the supplementary material.

##  3 Experiment

In this section, we conduct experiments to answer the following questions:

  * •

Can the proposed method enhance the performances of diverse off-policy MFRL
algorithms in various environments?

  * •

What attributes to the success of our method the most?

We first describe our experimental setup and show the main results against
relevant baselines to show that our method is generally applicable to any off-
policy MFRL methods with $Q$-learning and experience replays, and largely
improves their performance. Then we perform an ablation study of MQN to
analyze the most crucial components of it.

###  3.1 Experimental setup

Off-policy MFRL algorithms. We validate the effectiveness of our model-
augmented $Q$-learning with the following algorithms: Soft Actor-Critic (SAC)
(Haarnoja et al.,, 2018b), Twin Delayed Deep Deterministic (TD3), and Rainbow
(Hessel et al.,, 2018). Here, we apply data-efficient Rainbow (van Hasselt et
al.,, 2019) since its sample efficiency is dramatically higher than the
original and SimPLe. We modify these base algorithms to obtain different
variants of our method, namely Model-augmented SAC (MSAC), Model-augmented TD3
(MTD3), and Model-augmented Rainbow (MRainbow). We consider SAC, TD3, and
Rainbow because they and their variants are state-of-the-art off-policy RL
algorithms for continuous and discrete control tasks, respectively. We
emphasize that when applying our method to given algorithms, we do not alter
the original hyper-parameters to show that our method effectively improves the
base algorithm’s performance without any hyperparameter tuning. The detailed
configuration of the hyper-parameters are provided in the supplementary
material.

Environments. Although most RL algorithms frequently have used MuJoCo
environments (Todorov et al.,, 2012), they are not freely available to
everyone since they belong to a commercial physics engine. Accordingly, we
consider alternative free implementations of the original MuJoCo environments,
called PyBullet Gymperium222https://github.com/benelot/pybullet-gym, and other
free environments in the OpenAI Gym (Brockman et al.,, 2016a). We validate our
method on the following standard continuous control tasks:
HumanoidPyBulletEnv-0 (Humanoid), HalfCheetahPyBulletEnv-v0 (HalfCheetah),
HopperPyBulletEnv-v0 (Hopper), BipedalWalkerHardcore-v3, and Pendulum*. All of
environments are supported by PyBullet Gymperium except for the last two
tasks. BipedalWalkerHardcore-v3 belongs to Box2D continuous control tasks in
the OpenAI gym and Pendulum∗ is a sparse reward environment which is a
modification from Pendulum-v0 from the OpenAI Gym. In Pendulum∗, agents
receive the reward of $+1$ only if the rod is at the upright position
$(-3/\pi,3/\pi)$ in 100 continuous steps. In the case of discrete control
tasks, we validate our method on Atari games. We provide the details of the
environments we used for the experiments in the supplementary material.

|  |  |
---|---|---|---
(a) Humanoid (SAC) | (b) HalfCheetah (SAC) | (c) Hopper (SAC) | (d) BWH (SAC)
|  |  |
(e) Humanoid (TD3) | (f) HalfCheetah (TD3) | (g) Hopper (TD3) | (h) BWH (TD3)
Figure 2: Learning curves of off-policy RL algorithms SAC and TD3 on PyBullet
environments (Humanoid, HalfCheetah, and Hopper) and BWH. Here, MSAC-MPER and
MTD3-MPER, which represent variants of SAC and TD3 by our method,
respectively, outperform other RL algorithms. CDE and PB denote reward shaping
methods: curiosity-driven exploration (Pathak et al.,, 2017) and potential-
based reward shaping (Zou et al.,, 2019), respectively. The solid line and
shaded regions represent the mean and standard deviation, respectively, across
five runs with random seeds.

Reward shaping methods. We compare our MReward against the following reward
shaping methods.

  * •

Curiosity-driven exploration (CDE) (Pathak et al.,, 2017): This method uses
additional intrinsic rewards to promote exploration:

$r_{t}^{i}=\frac{\eta}{2}\|\widehat{\phi}\left(s_{t+1}\right)-\phi\left(s_{t+1}\right)\|,$

where $\phi(s_{t})$ is a feature vector and $\widehat{\phi}(s_{t})$ is a
predicted feature vector for $s_{t}$, respectively.

  * •

Potential-based method (PB) (Ng et al.,, 1999; Wiewiora et al.,, 2003; Devlin
& Kudenko,, 2012): This method uses a modified reward $R_{t}=r_{t}+F$, where

$F_{t}(s,a,s^{\prime},a^{\prime})=-\gamma\Psi(s^{\prime},a^{\prime})+\Psi(s,a)$

and $\Psi$ is a real-valued function. Although there exist various methods for
choosing $\Psi$, we use the value network as $\Psi$ following (Zou et al.,,
2019), as we do not assume any prior knowledge about the given task.

Sampling methods. We compare the performance of our MPER against the following
experience replays.

  * •

Experience Replay with Uniform Sampling at Random (RANDOM) : Sampling
transitions uniformly at random.

  * •

Priotized Experience Replay (PER) (Schaul et al.,, 2016): Rule-based
prioritized sampling of the transitions based on TD-errors.

  * •

Experience Replay with Likelihood-free Importance Weights (LFIW) (Sinha et
al.,, 2020): Learning-based sampling method which predicts the importance of
each experience. Since the authors only validated it on continuous control
tasks with SAC and TD3, we compare against it in the same settings.

|  |  |
---|---|---|---
(a) Amidar | (b) Assault | (c) BankHeist | (d) DemonAttack
|  |  |
(e) Frostbite | (f) Jamesbond | (g) Krull | (h) Seaquest
Figure 3: Learning curves of Rainbow on Atari games. Here, MRainbow-MPER is a variant of Rainbow, where all of MQN, MPER, and MReward are applied. Although Rainbow basically adopts PER, we denote it as Rainbow-PER for consistency. One can observe that MRainbow-MPER overwhelms other methods. The solid line and shaded regions represent the mean and standard deviation, respectively, across five runs with random seeds.  |
---|---
(a) Pendulum∗ (SAC) | (b) Pendulum∗ (TD3)
Figure 4: Learning curves of off-policy RL algorithms SAC and TD3 on
Pendulum∗. Here, MSAC-MPER and MTD3-MPER, which represent variants of SAC and
TD3 by our method, respectively, outperform other RL algorithms. CDE and PB
denote reward shaping methods: curiosity-driven exploration (Pathak et al.,,
2017) and potential-based reward shaping (Zou et al.,, 2019), respectively.
The solid line and shaded regions represent the mean and standard deviation,
respectively, across five runs with random seeds.

###  3.2 Main Experimental Results

We consider three types of environments for our main experiments: Pybullet,
sparse Pendulumn, and Atari games.

Non-sparse reward environment. Figure A.2.(a)-(d) and Figure A.2.(e)-(h) show
the learning curves of SAC and TD3 on PyBullet environments and
BipdalWalkerHardcore-v3, respectively. The sampling method LFIW drives better
results compared to SAC and TD3 with PER and RANDOM, but its variance in
performance is much larger, which means that the training is less stable. On
the other hand, our MAQ consistently outperforms baselines in all tested
cases. In particular, It significantly improves the performance of all off-
policy RL algorithms on the environment with high-dimensional state and action
spaces, i.e., HumanoidPyBulletEnv-v0. This impressive performance gain is made
possible by two factors. 1) Our representation for the state-action pair
$(s,a)$, which is the input for the $Q$-function, improves by parallel
unification of MQN, MReward, and MPER. 2) The Model-augmented TD error in Eq.
13 is improved, as the model estimation for the reward and the transition
improves. 3) MPER samples the transitions that allows for further improvements
in the accurate estimation of the Model-augmented TD error in Eq. 13.

Atari game environments. Next, since we focus on how to improve $Q$-learning,
we validate our method with Rainbow (van Hasselt et al.,, 2019), which is a
state-of-the-art $Q$-learning framework without a policy network. Figure 3
shows the learning curves of data-efficient Rainbow (van Hasselt et al.,,
2019) on various Atari games. MRainbow in the figure is a variant of Rainbow
modified with our MAQ, that consists of MQN, MPER, and MReward. We can observe
that MRainbow achieves overwhelmingly larger performance compared with the
base Rainbow, with both PER and RANDOM on various Atari games. We provide
additional Atari results in the supplementary material.

Sparse reward environment. Finally, we compare ours on the sparse reward
environment, i.e., Pendulum∗ against two reward shaping methods, CDE, PE in
Figure 4. Similarly to the experimental results on environments with dense
rewards, MSAC-MPER and MTD3-MPER, which utilize our method, achieves
overwhelmingly better performance over the baseline algorithms. While our MQN
succesfully trained on all five instances, most baseline algorithms and
sampling methods were unable to learn how to get high cumulative rewards.
Although CDE was able to successfully learn good policies on certain
instances, it failed on others.

###  3.3 Ablation Study

We now analyze what components in our method are crucial to its improvement of
MFRL algorithms’ performance.

Effectiveness of Sampling Methods. First, we verify the effectivness of MPER.
Since this method is only available with MQN, we set the baseline as MSAC and
compare our sampling method MPER with RANDOM, PER, LFIW under MSAC and MTD3,
respectively. Figure 5.(a)-(b) shows the learning curves of MSAC with
different sampling methods. We can observe that LFIW and PER are inferior to
RANDOM, even with the MQN $\mathcal{M}_{\theta}$. The main reason is that
these two methods focus on sampling transitions that are beneficial in
updating the $Q$-network only (i.e., TD-errors and importance weights for the
$Q$-loss). Since estimation of the $Q$-value is only a single component of
MQN, this is suboptimal in improving our Mode-augmented TD error objective,
and thus we need MPER which seeks to improve the estimation of the reward,
state, and $Q$-values.

Effectiveness of model learning. Second, we analyze the effectiveness of the
model estimation in MAQ. To verify it, we examine the performance of the MSAC
variants without the reward or transition estimators. We refer to them as
reward-considering (RSAC) with reward-considering PER (RPER) and transition-
considering (TSAC) with transition-considering PER (TPER) respectively. Figure
5.(c)-(d) show the learning curves of each method. SAC-PER performs the worst
and adding each component leads to performance improvement. Also, TSAC-TPER
outperforms RSAC-RPER, which suggests that learning the transition map may be
more beneficial in improving the Q-value estimation or the representational
power of the state-action representations.

Effectiveness of Increased Parameters. One may suspect that if MSAC achieves
improved performance due to the increased size of the $Q$-network, with
additional estimators. To show that this is not the case, we use a base SAC
with its hidden layers dramatically increased, from $(300,300)$ to
$(2500,2500)$ for both the Q-network and the policy network. Figure A.5.(a)
shows the effect of changes by learning curves. We observe that increasing the
hidden layer size alone does not improve the performance and may even lead to
the performance degeneration. From these results, it is clear that the
improvements of our method is not simply coming from the increased size of
hidden layers.

Effectiveness of interaction and parameter sharing between the estimators. We
further verify the effectiveness of parameter sharing and interaction across
the model and the $Q$-value estimators in MQN, to see where the improvements
come from. For this experience, we consider a MQN with separate networks for
model and $Q$-value estimation (without parameter sharing) in Figure A.5.(b),
and a MQN that utilizes the real reward instead of MReward in (13) to remove
the interactions across estimators in Figure A.5.(b). One can observe that
removing the parameter sharing critically degrades the sample efficiency of
MSAC-MPER. Moreover, although MSAC without the interaction learns well, MSAC-
MPER with the interaction finally overtakes it, which shows that learning with
the estimator interactions is also useful.

Ananalys of value estimation errors. To check if our proposed method
effectively deals with the overestimation problem, we check how close are the
estimated $Q$-values to the true returns. To this end, we compute $Q$-values
of initial states and computed true returns via five instances in Figure
A.5.(c). We can observe that MSAC-MPER’s $Q$-estimation converges to the real
returns until less than 0.1M steps. On the contrary, SAC-RANDOM’s
$Q$-estimation is largely overestimated and does not converge to the true
returns before 0.5M steps. These results show that our method dramatically
improves the precision of $Q$-learning.

Analysis of the relation between model- and TD-errors. Figure A.5.(d) shows
that model- and TD-errors are highly related, with our MSAC-MPER. We observe
that small model-errors generally lead to small TD-errors estimated with the
true reward (Eq. 2). Notice that at the beginning, as model-errors decreases
rapidly, TD-errors also decreases sharply in succession. This is a direct
evidence that estimation of the model with our MQL objective (13) is indeed
helpful with the $Q$-learning itself. However, in the case of SAC-RANDOM,
which uses a separate model network without the interaction term, the TD-
errors and model estimation errors are uncorrelated, and TD-error decreases
much slowly compared to MSAC-MPER.

Hopper |  HalfCheetah |  Hopper |  HalfCheetah
---|---|---|---
|  |  |
(a) Sampling effect | (b) Sampling effect | (c) Model’s effect | (d) Model’s effect
Figure 5: (a/b): Learning curves of MSAC with different sampling methods. Under MSAC, we observe that the proposed MPER outperforms other sampling methods. (c/d): Learning curves for different variants of MPER for the ablation study. TSAC-TPER refers a variant of MPER with only the transition estimator that prioritizes the PER only with the transition estimation error, and RSAC-RPER refers to a version that only considers and estimates reward. One can observe that learning of both reward and transition maps is the most effective. The solid lines and shaded regions represent the mean and standard deviations across five runs with random seeds.  Humanoid |  Humanoid |  Humanoid |  Humanoid
---|---|---|---
|  |  |
(a) Network size effect | (b) Interaction effect | (c) Estimation bias | (d) Model- and TD-errors
Figure 6: (a): Learning curves of SAC with larger networks. Here, large actor
and critic are networks whose hidden layer contains $(2500,2500)$ hidden
units, respectively. This shows that simply increasing the number of
parameters in the value and the policy network leads to marginal or worse
performance. (b): Measured estimation bias in the $Q$-values for MSAC-MPER and
SAC-RANDOM, respectively. We observe that MSAC-MPER quickly converges to the
true returns, with less than 0.1M steps. On the contrary, SAC-RANDOM largely
overestimate the value in earlier training steps. (c): Comparison with and
without the interaction of the model and value estimator in (13) or the
parameter sharing between the estimators of MQN. MSAC-MPER without interaction
leverages the true reward to train the model, and MSAC-MPER without parameter
sharing leverages separate networks for each estimator. (d): Normalized model-
and TD-error curves. We can observe that the model- and TD-errors have the
same tendency under MSAC-MPER unlike SAC-RANDOM with a separate model-
estimation network, which is a direct evidence that the model estimation in
MQL is indeed helpful in reducing the TD-errors. The solid lines and shaded
regions represent the mean and standard deviations across five runs with
random seeds.

##  4 Conclusion

We proposed model-augmented $Q$-learning (MQL), which utilizes a model-
augmented $Q$-network (MQN) that estimates not only the $Q$-value but also the
reward and transition maps via parameter sharing. However, instead of simply
relying on shared parameters, we promote interactions between MQN’s estimators
by deriving the equation (13) given a model-augmented reward (MReward). We
also proved that MQL is policy-invariance, i.e., it does not affect the set of
Pareto-optimal policies. We further propose Model-augmented Prioritized
Experience Replay (MPER), which samples the past experiences based on the
model estimation errors, as well as the TD errors.

The advantages of our proposed method are as follows. First, It simple to
implement and is generally applicable to any MFRL algorithms that utilize
$Q$-networks. Second, MQL dramatically increases sample efficiency of state-
of-the-art MFRL algorithms: SAC, TD3, and Rainbow. Third, it largely
alleviates the understimation or overestimation of the value with the
conventional $Q$-learning. Fourth, the computational cost of our algorithms is
almost the same as the original algorithms. Finally, MQL is effective in
various environments, including ones with the sparse rewards.

## References

  * Agarwal et al., (2018) Agarwal, Praveen, Jleli, Mohamed, & Samet, Bessem. 2018.  Banach Contraction Principle and Applications.  Pages 1–23 of: Fixed Point Theory in Metric Spaces.  Springer.
  * Brittain et al., (2019) Brittain, Marc, Bertram, Josh, Yang, Xuxi, & Wei, Peng. 2019.  Prioritized Sequence Experience Replay.  arXiv preprint arXiv:1905.12726.
  * Brockman et al., (2016a) Brockman, Greg, Cheung, Vicki, Pettersson, Ludwig, Schneider, Jonas, Schulman, John, Tang, Jie, & Zaremba, Wojciech. 2016a.  Openai gym.  arXiv preprint arXiv:1606.01540.
  * Brockman et al., (2016b) Brockman, Greg, Cheung, Vicki, Pettersson, Ludwig, Schneider, Jonas, Schulman, John, Tang, Jie, & Zaremba, Wojciech. 2016b.  OpenAI Gym.
  * Clavera et al., (2018) Clavera, Ignasi, Rothfuss, Jonas, Schulman, John, Fujita, Yasuhiro, Asfour, Tamim, & Abbeel, Pieter. 2018.  Model-Based Reinforcement Learning via Meta-Policy Optimization.  Pages 617–629 of: Conference on Robot Learning.
  * Devlin & Kudenko, (2012) Devlin, Sam Michael, & Kudenko, Daniel. 2012.  Dynamic potential-based reward shaping.  Pages 433–440 of: Proceedings of the 11th International Conference on Autonomous Agents and Multiagent Systems.  IFAAMAS.
  * Fujimoto et al., (2018) Fujimoto, Scott, van Hoof, Herke, & Meger, David. 2018.  Addressing function approximation error in actor-critic methods.  In: ICML.
  * Haarnoja et al., (2018a) Haarnoja, Tuomas, Zhou, Aurick, Hartikainen, Kristian, Tucker, George, Ha, Sehoon, Tan, Jie, Kumar, Vikash, Zhu, Henry, Gupta, Abhishek, Abbeel, Pieter, et al. 2018a.  Soft actor-critic algorithms and applications.  arXiv preprint arXiv:1812.05905.
  * Haarnoja et al., (2018b) Haarnoja, Tuomas, Zhou, Aurick, Abbeel, Pieter, & Levine, Sergey. 2018b.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In: ICML.
  * Harutyunyan et al., (2015) Harutyunyan, Anna, Devlin, Sam, Vrancx, Peter, & Nowé, Ann. 2015.  Expressing arbitrary reward functions as potential-based advice.  In: Proceedings of the AAAI Conference on Artificial Intelligence, vol. 29.
  * Hasselt, (2010) Hasselt, Hado V. 2010.  Double Q-learning.  In: NeurIPS.
  * Hessel et al., (2018) Hessel, Matteo, Modayil, Joseph, Van Hasselt, Hado, Schaul, Tom, Ostrovski, Georg, Dabney, Will, Horgan, Dan, Piot, Bilal, Azar, Mohammad, & Silver, David. 2018.  Rainbow: Combining improvements in deep reinforcement learning.  In: AAAI.
  * Hou et al., (2017) Hou, Yuenan, Liu, Lifeng, Wei, Qing, Xu, Xudong, & Chen, Chunlin. 2017.  A novel ddpg method with prioritized experience replay.  In: SMC.
  * Hu et al., (2020) Hu, Yujing, Wang, Weixun, Jia, Hangtian, Wang, Yixiang, Chen, Yingfeng, Hao, Jianye, Wu, Feng, & Fan, Changjie. 2020.  Learning to Utilize Shaping Rewards: A New Approach of Reward Shaping.  Advances in Neural Information Processing Systems, 33.
  * Janner et al., (2019) Janner, Michael, Fu, Justin, Zhang, Marvin, & Levine, Sergey. 2019.  When to trust your model: Model-based policy optimization.  Pages 12519–12530 of: Advances in Neural Information Processing Systems.
  * Kaiser et al., (2019) Kaiser, Lukasz, Babaeizadeh, Mohammad, Milos, Piotr, Osinski, Blazej, Campbell, Roy H, Czechowski, Konrad, Erhan, Dumitru, Finn, Chelsea, Kozakowski, Piotr, Levine, Sergey, et al. 2019\.  Model-based reinforcement learning for atari.  arXiv preprint arXiv:1903.00374.
  * Kingma & Ba, (2014) Kingma, Diederik P, & Ba, Jimmy. 2014.  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980.
  * Konda & Tsitsiklis, (2000) Konda, Vijay R, & Tsitsiklis, John N. 2000.  Actor-critic algorithms.  Pages 1008–1014 of: Advances in neural information processing systems.  Citeseer.
  * Kurutach et al., (2018) Kurutach, Thanard, Clavera, Ignasi, Duan, Yan, Tamar, Aviv, & Abbeel, Pieter. 2018\.  Model-Ensemble Trust-Region Policy Optimization.  In: International Conference on Learning Representations.
  * Lan et al., (2019) Lan, Qingfeng, Pan, Yangchen, Fyshe, Alona, & White, Martha. 2019.  Maxmin Q-learning: Controlling the Estimation Bias of Q-learning.  In: International Conference on Learning Representations.
  * Liang & Zhang, (2020) Liang, Sicong, & Zhang, Yu. 2020.  A Simple General Approach to Balance Task Difficulty in Multi-Task Learning.  arXiv preprint arXiv:2002.04792.
  * Luo et al., (2018) Luo, Yuping, Xu, Huazhe, Li, Yuanzhi, Tian, Yuandong, Darrell, Trevor, & Ma, Tengyu. 2018.  Algorithmic Framework for Model-based Deep Reinforcement Learning with Theoretical Guarantees.  In: International Conference on Learning Representations.
  * Mnih et al., (2015) Mnih, Volodymyr, Kavukcuoglu, Koray, Silver, David, Rusu, Andrei A, Veness, Joel, Bellemare, Marc G, Graves, Alex, Riedmiller, Martin, Fidjeland, Andreas K, Ostrovski, Georg, et al. 2015\.  Human-level control through deep reinforcement learning.  Nature, 518(7540), 529.
  * Nagabandi et al., (2018) Nagabandi, Anusha, Kahn, Gregory, Fearing, Ronald S, & Levine, Sergey. 2018.  Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning.  Pages 7559–7566 of: 2018 IEEE International Conference on Robotics and Automation (ICRA).  IEEE.
  * Ng et al., (1999) Ng, Andrew Y, Harada, Daishi, & Russell, Stuart. 1999.  Policy invariance under reward transformations: Theory and application to reward shaping.  Pages 278–287 of: ICML, vol. 99.
  * Pathak et al., (2017) Pathak, Deepak, Agrawal, Pulkit, Efros, Alexei A, & Darrell, Trevor. 2017.  Curiosity-driven exploration by self-supervised prediction.  Pages 16–17 of: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops.
  * Rajeswaran et al., (2020) Rajeswaran, Aravind, Mordatch, Igor, & Kumar, Vikash. 2020.  A Game Theoretic Framework for Model Based Reinforcement Learning.  arXiv, arXiv–2004.
  * Schaul et al., (2016) Schaul, Tom, Quan, John, Antonoglou, Ioannis, & Silver, David. 2016.  Prioritized experience replay.  In: ICLR.
  * Schulman et al., (2015) Schulman, John, Levine, Sergey, Abbeel, Pieter, Jordan, Michael, & Moritz, Philipp. 2015.  Trust region policy optimization.  Pages 1889–1897 of: International conference on machine learning.
  * Schulman et al., (2017) Schulman, John, Wolski, Filip, Dhariwal, Prafulla, Radford, Alec, & Klimov, Oleg. 2017.  Proximal policy optimization algorithms.  arXiv preprint arXiv:1707.06347.
  * Sinha et al., (2020) Sinha, Samarth, Song, Jiaming, Garg, Animesh, & Ermon, Stefano. 2020.  Experience Replay with Likelihood-free Importance Weights.  arXiv preprint arXiv:2006.13169.
  * Sutton, (1991) Sutton, Richard S. 1991.  Dyna, an integrated architecture for learning, planning, and reacting.  ACM Sigart Bulletin, 2(4), 160–163.
  * Todorov et al., (2012) Todorov, Emanuel, Erez, Tom, & Tassa, Yuval. 2012.  Mujoco: A physics engine for model-based control.  In: IROS.
  * van Hasselt et al., (2019) van Hasselt, Hado P, Hessel, Matteo, & Aslanides, John. 2019.  When to use parametric models in reinforcement learning?  Pages 14322–14333 of: Advances in Neural Information Processing Systems.
  * Vinyals et al., (2019) Vinyals, Oriol, Babuschkin, Igor, Czarnecki, Wojciech M, Mathieu, Michaël, Dudzik, Andrew, Chung, Junyoung, Choi, David H, Powell, Richard, Ewalds, Timo, Georgiev, Petko, et al. 2019\.  Grandmaster level in StarCraft II using multi-agent reinforcement learning.  Nature, 575(7782), 350–354.
  * Wang & Ross, (2019) Wang, Che, & Ross, Keith. 2019.  Boosting Soft Actor-Critic: Emphasizing Recent Experience without Forgetting the Past.  arXiv preprint arXiv:1906.04009.
  * Wiewiora et al., (2003) Wiewiora, Eric, Cottrell, Garrison W, & Elkan, Charles. 2003.  Principled methods for advising reinforcement learning agents.  Pages 792–799 of: Proceedings of the 20th International Conference on Machine Learning (ICML-03).
  * Zha et al., (2019) Zha, Daochen, Lai, Kwei-Herng, Zhou, Kaixiong, & Hu, Xia. 2019.  Experience replay optimization.  In: IJCAI.
  * Zou et al., (2019) Zou, Haosheng, Ren, Tongzheng, Yan, Dong, Su, Hang, & Zhu, Jun. 2019.  Reward shaping via meta-learning.  arXiv preprint arXiv:1901.09330.

Supplementary Material:

Model-Augmented $Q$-learning

##  Appendix A Environment Description

###  A.1 PyBullet Environments

PyBullet environments are the open-source implementations of the OpenAI Gym
MuJoCo (Todorov et al.,, 2012) environments, which is currently one of the
most widely used toolkits for developing and comparing reinforcement learning
algorithms. Since MuJoCo environments are commercial, they hinder open
research, so we do not use MuJoCo environments in the manuscript. Besides,
these environments are known to be more difficult due to the more realistic
improvement, e.g., adding energy cost, of the existing MuJoCo environments. We
provide PyBullet tasks applied in the manuscript:

HumanoidPyBulletEnv-v0 is an environment to control a three-dimensional
bipedal robot to walk quickly without falling over.

HalfCheetahPyBulletEnv-v0 is an environment to control a two-dimensional
cheetah robot for learning sprint.

HopperPyBulletEnv-v0 is an environment to control a two-dimensional one-legged
robot hop-forward quickly without falling over.

|  |
---|---|---
(a) HumanoidPyBulletEnv-v0 | (b) HalfCheetahPyBulletEnv-v0 | (c) HopperPyBulletEnv-v0
Figure A.1: Pybullet environments.

###  A.2 OpenAI Gym Environments

OpenAI Gym (Brockman et al.,, 2016b) supports continuous control environments
that belong to classic or Box2D simulators. We conduct experiments on the
following environments among them.

Pendulum∗ is an environment which objective is to balance a rod in the upright
position as long as possible. It is a variant of Pendulum-v0 that is supported
by OpenAI gym. To make it sparser, we impose the following condition: The
pendulum begins to receive +1 reward if maintaining the rod in the upright
position more than 100 steps continuously.

BipedalWalkerHardcore-v3 is an environment to control a robot, a variant of
BipedalWalker-v3 in OpenAI gym. The robot’s objective is to move forward as
far as possible while solving many obstacles.

|
---|---
(a) BipedalWalkerHardcore-v3 | (b) Pendulum-v0
Figure A.2: OpenAi Gym environments.

Table A.1 shows the observation and action spaces and the maximum steps for
each episode (horizon) in Pybullet and OpenAI gym environments that we
considered. Here, $\mathbb{R}$ and $[-1,1]$ denote sets of real numbers and
those between $0$ and $1$, respectively.

Environment | Observation space | Action space | Horizon
---|---|---|---
HumanoidPybulletEnv-v0 | $\mathbb{R}^{{}^{44}}$ | $\left[-1,1\right]^{17}$ | $1000$
HalfCheetahPybulletEnv-v0 | $\mathbb{R}^{{}^{26}}$ | $\left[-1,1\right]^{6}$ | $1000$
HopperCheetahPybulletEnv-v0 | $\mathbb{R}^{{}^{15}}$ | $\left[-1,1\right]^{3}$ | $1000$
BipedalWalkerHardcore-v3 | $\mathbb{R}^{{}^{24}}$ | $\left[-1,1\right]^{4}$ | $2000$
Pendulum∗ | $\mathbb{R}^{{}^{3}}$ | $\left[-1,1\right]^{1}$ | 200
Table A.1: Dimensions of observation and action spaces for continuous control
environments

###  A.3 Discrete Control Environments

We explain Atari environments that we considered in the manuscript and this
supplementary material. The objective of RL agents is to learn a policy by
observing the screen (RGB) to get high cumulative rewards.

Alien: An environment where player should destroy all alien eggs in the RGB
screen with escaping aliens.

Amidar: An environment similar to MsPacman. In this environment, agents
control a monkey in a fixed rectilinear lattice to eat pellets as much as
possible while avoiding chasing masters.

Assault: An environment to control a spaceship. which objective is to
eliminate the enemies.

Asterix: An environment to control a tornado. Its objective is to eat
hamburgers in the screen with avoiding dynamites.

|  |  |
---|---|---|---
(a) Alien | (b) Amidar | (c) Assault | (d) Asterix
Figure A.3: Atari Games: Alien, Amidar, Assault, and Asterix.

BankHeist: An environment to control a robber. The objective of the game is to
rob banks as many as possible while avoiding the police in maze-like cities..

DemonAttack: An environment to control a guardian.The guardian should kill
demons that attacks from above.

Frostbite: An environment to control a man who should collect ice blocks to
make his igloo. His objective is to collect 15 ice blocks while avoiding some
opponents, e.g., crabs and birds.

Gopher: An environment to control a farmer. The farmer should protect three
carrots from a gopher.

|  |  |
---|---|---|---
(a) BankHeist | (b) DemonAttack | (c) FrostBite | (d) Gopher
Figure A.4: Atari Games: BankHesit, DemonAttack, Frostbite, and Gopher.

Jamesbond: An environment to control a vehicle. The objective is to move
forward while avoiding and attacking enemies.

Kangaroo: An environment to control a mother kangaroo. The kangaroo’s
objective is to rescue her son while climbing.

Krull: An environment to control a player following a file of the same title.
The player should complete stages, which are main parts in the film.

Seaquest: An environment to control a submarine. Its objective is to rescue
divers while attacking enemies by missiles.

|  |  |
---|---|---|---
(a) Jamesbond | (b) Kangaroo | (c) Krull | (d) Seaquest
Figure A.5: Atari Games: Jamesbond, Kangaroo, Krull, and Seaquest.

##  Appendix B Training details

Table B.1 shows hyper-parameters that we used in experiments of the
manuscript. Notice that MQN barely needs hyper-parameters. In the case of
Rainbow, instead of taking RGB array as input and output, MQN’s transition map
takes a context by CNN layers for the current and next RGB arrays as input and
output.

Parameter | Value
---|---
Shared |
Batch size (continuous control environments) |  $256$ (SAC), $100$ (TD3), $32$ (Rainbow)
Buffer size | $10^{6}$
Target smoothing coefficient $(\tau)$ for soft update | $5\times 10^{-3}$
Initial prioritized experience replay buffer exponents $(\alpha,\beta)$ 333$\beta$ increases to $1.0$ by the rule $\beta=0.4\eta+1.0(1-\eta)$, where $\eta=$ the current step/the maximum steps. |  $(0.7,$ $0.4)$ (SAC/TD3), $(0.5,$ $0.4)$ (Rainbow)
Discount factor for the agent reward ($\gamma$) |  $0.98$ (SAC/TD3), $0.99$ (Rainbow)
Number of initial random actions (continuous control environments) | $5\times 10^{3}$
Number of initial random actions (discrete control environments) |  $10,000$ (SAC/TD3) $1,600$ (Rainbow)
Optimizer |  Adam (Kingma & Ba,, 2014)
Nonlinearity | ReLU
Replay period | 64 (SAC/TD3), 1 (Rainbow)
Gradient step | 64 (SAC/TD3), 1 (Rainbow)
MQN |
MReward coefficients $(\zeta_{1}$, $\zeta_{2})$ |  $(10^{-3}$, $10^{-3})$
Likelihood-free Importance Weights |
Temperature $(T)$ | 5
Hidden units per layer | 256, 256
Fast replay buffer size | $10^{4}$
TD3 |
Hidden units per layer | 400, 300
Learning rate | $10^{-3}$
Policy update frequency | $2$
Gaussian action and target noises |  $0.1,$ $0.2$
Target noise clip | 0.5
Target network update |  Soft update with interval $1$
SAC |
Hidden units per layer | 400, 300
Learning rate | $7.3\times 10^{-4}$
Target entropy |  $-\dim A$ ($A$ is action space)
Soft update with interval | $1$
Rainbow |
Observation down-sampling for Atari RGB |  $84\times 84$ with grey-scaling
CNN channels for Atari environments | 32, 64
CNN filter size for Atari environments |  $5\times 5$, $5\times 5$
CNN stride for Atari environments | 5, 5
Action repetitions and Frame stack | 4
Reward clipping |  True ($[-1,1]$)
Terminal on loss of life | True
Max frames per episode | $1.08\times 10^{5}$
Target network update | Hard update (every 2,000 updates)
Support of $Q$-distribution | 51
$\epsilon$ for Adam optimizer | $1.5\times 10^{4}$
Learning Rate | $10^{-4}$
Max gradient norm | 10
Noisy nets parameter | 0.1
Multi-step return length | 20
$Q$-network’s hidden units per layer | 256
Table B.1: Hyper-parameters

##  Appendix C The Proof of Theorem 1

###### Proof.

We first define the following notations

$\displaystyle\|(x_{1},\cdots,x_{n})\|_{2}^{2}$
$\displaystyle=\frac{1}{2}\sum_{i=1}^{n}x_{i}^{2},$
$\displaystyle\|(x_{1},x_{2},x_{3},\cdots,x_{n})||$
$\displaystyle=||x_{1}||_{\infty(S)}+\sum_{i=1}^{n}||x_{i}||_{2}^{2},$
$\displaystyle\|x||_{\infty(S)}$ $\displaystyle=\max_{s\in S}x(s),$

where $S$ is a state space. Let $\boldsymbol{R}$ and $\boldsymbol{S}$ be
defined by

$\displaystyle\boldsymbol{\boldsymbol{R}}$
$\displaystyle=\\{\mathcal{R}(s,\pi(s))|s\in S\\}\in\mathbb{R}^{\dim(S)},$
$\displaystyle\boldsymbol{\boldsymbol{S}}$
$\displaystyle=\\{\mathcal{T}(s,\pi(s))|s\in S\\}\in\mathbb{R}^{\dim(S)},$

where $\left(\mathcal{R},\mathcal{T}\right)$ is an environment’s model. By
defining an operator $\mathcal{P}^{\pi}$ by

$\displaystyle\mathcal{P}^{\pi}(q,r,s)$
$\displaystyle=\left(\mathcal{P}_{1}^{\pi}(q,r,s),\mathcal{P}_{2}^{\pi}(q,r,s),\mathcal{P}_{3}^{\pi}(q,r,s)\right),$
$\displaystyle\mathcal{P}_{1}^{\pi}(q,r,s)$
$\displaystyle=r+\zeta_{1}\|r-\boldsymbol{\boldsymbol{R}}\|_{2}+\zeta_{2}\|s-\boldsymbol{S}\|_{2}$
$\displaystyle+\gamma
q(\boldsymbol{S},\pi(\cdot|\boldsymbol{\boldsymbol{S}})),$
$\displaystyle\mathcal{P}_{2}^{\pi}(q,r,s)$
$\displaystyle=r-\kappa_{1}(r-\boldsymbol{\boldsymbol{R}}),$
$\displaystyle\mathcal{P}_{3}^{\pi}(q,r,s)$
$\displaystyle=s-\kappa_{2}(s-\boldsymbol{\boldsymbol{S}}),$

where $\kappa_{1},\kappa_{2}$ are small positive constants, $q\in
F:=\\{f|f:S\times
A\rightarrow\mathbb{R}\\},r\in\mathbb{R},s\in\mathbb{R}^{\dim(S)}$, one can
compute

$\displaystyle\left|\mathcal{P}_{2}^{\pi}(q^{1},r^{1},s^{1})-\mathcal{P}_{2}^{\pi}(q^{2},r^{2},s^{2})\right|$
$\displaystyle\leq\left|(1-\kappa_{1})(r^{1}-r^{2})\right|,$
$\displaystyle\left|\mathcal{P}_{3}^{\pi}(q^{1},r^{1},s^{1})-\mathcal{P}_{3}^{\pi}(q^{2},r^{2},s^{2})\right|$
$\displaystyle\leq\left|(1-\kappa_{2})(s^{1}-s^{2})\right|.$

Notice that the following inequalities are obtained:

$\displaystyle\left|\|x||_{2}-\|y||_{2}\right|$
$\displaystyle\leq\|x-y||_{2},$ (17) $\displaystyle\|x+y||_{2}$
$\displaystyle\leq\|x||_{2}+\|y||_{2}$ (18)

Using (17)-(18), one can also calculate

$\displaystyle||\mathcal{P}^{\pi}(q^{1},r^{1},s^{1})-\mathcal{P}^{\pi}(q^{2},r^{2},s^{2})||$
$\displaystyle\leq\left|\zeta_{1}\|r^{1}-r^{2}\|_{2}+\zeta_{2}\|s^{1}-s^{2}\|_{2}\right|$
$\displaystyle+\gamma\max_{s\in S}\left|q^{1}(s,\pi(\cdot|s))-\gamma
q^{2}(s,\pi(\cdot|s)\right|,$

so that

$\displaystyle||\mathcal{P}^{\pi}(q^{1},r^{1},s^{1})-\mathcal{P}^{\pi}(q^{2},r^{2},s^{2})||$
$\displaystyle\leq\max\left(\gamma,\zeta_{1},\zeta_{2},(1-\kappa_{1}),(1-\kappa_{2})\right)$
$\displaystyle\times||q^{1}-q^{2},r^{1}-r^{2},s^{1}-s^{2}||.$ (19)

Therefore, the Banach fixed point theorem (Agarwal et al.,, 2018) is available
by (19) due to

$\max\left(\gamma,\zeta_{1},\zeta_{2},(1-\kappa_{1}),(1-\kappa_{2})\right)<1,$

so there exists a unique $\left(q^{*},r^{*},s^{*}\right)$ such that

$\mathcal{P}^{\pi}\left(q^{*},r^{*},s^{*}\right)=\left(q^{*},r^{*},s^{*}\right).$

Finally, let $q^{0}\in F$ such that

$q^{0}=\boldsymbol{r}+\gamma
q^{0}(\boldsymbol{S},\pi(\cdot|\boldsymbol{\boldsymbol{S}})).$

Since

$\displaystyle r^{*}$
$\displaystyle=r^{*}-\kappa_{1}(r^{*}-\boldsymbol{\boldsymbol{R}}),$
$\displaystyle s^{*}$
$\displaystyle=s^{*}-\kappa_{2}(s^{*}-\boldsymbol{\boldsymbol{S}}),$

it is easy to get
$r^{*}=\boldsymbol{\boldsymbol{R}},s^{*}=\boldsymbol{\boldsymbol{S}}.$
Therefore,

$q^{*}=\boldsymbol{r}+\gamma
q^{*}(\boldsymbol{S},\pi(\cdot|\boldsymbol{\boldsymbol{S}})).$

Since $q^{*}$ is a unique solution in $F,$ one can conclude that $q^{*}\equiv
q_{0}$. This completes the proof. ∎

##  Appendix D Related Work

Reward shaping. Our reward estimator in the MQN is closely related to the
reward shaping methods. The most well-known and frequently used methods for
reward shaping are the potential-based reward shaping (PBRS) (Ng et al.,,
1999) and its variants (Wiewiora et al.,, 2003; Devlin & Kudenko,, 2012;
Harutyunyan et al.,, 2015). Ng et al., (1999); Wiewiora et al., (2003); Devlin
& Kudenko, (2012) prove that the policy-invariance is guaranteed even if the
reward is modified by potential functions that depend only on states. Since
reward shaping usually requires heavy prior knowledge of the given tasks, Hu
et al., (2020) proposed to learn the weight of a given potential function, and
Zou et al., (2019) used meta-learning to learn the reward estimator.

Model-based RL. Although our approach is different from model-based RL
algorithms, we briefly introduce some of them. There are different types of
Model-based RL (MBRL), and the most popular approach is Dyna-style (Sutton,,
1991), in which algorithms generate fictitious experiences to train agents.
Although there exists various methods (Kurutach et al.,, 2018; Luo et al.,,
2018; Clavera et al.,, 2018; Janner et al.,, 2019; Rajeswaran et al.,, 2020)
for MBRL, the common strategy is to first learn the environment model and use
it to generate fictitious experiences for learning an agent’s policy. Due to
its ability to generate transitions, MBRL’s sample efficiency is remarkable on
certain tasks, but they require prior knowledge of environments such as a
termination function on state and reward function, for a given state-action
pair. Furthermore, MBRL requires much larger computing costs over MFRL
algorithms.

Off-policy model-free RL. In the case of $Q$-learning without policy newtorks,
Rainbow (Hessel et al.,, 2018) that combines various techniques to extend the
original DQN learning (Mnih et al.,, 2015) is one of the state-of-the-art
methods in Atari game environments. Although there is a model-based approach
for Atari games (Kaiser et al.,, 2019), data-efficient Rainbow (van Hasselt et
al.,, 2019) that we use as a baseline outperforms it. In the case of the
actor-critic architecture, twin delayed DDPG (TD3) (Fujimoto et al.,, 2018)
and soft actor-critic (SAC) (Haarnoja et al.,, 2018b) are frequently used the
state-of-the art methods. TD3 employs double $Q$-networks, target policy
smoothing, and different frequencies to update a policy and $Q$-networks, to
reduce overestimation bias. SAC also adopts double $Q$-learning and utilize
the entropy measure of an agent policy to the reward to encourage the
exploration of the agent.

Experience replay. Prioritized experience replay (PER) (Hessel et al.,, 2018;
Hou et al.,, 2017; Schaul et al.,, 2016; Wang & Ross,, 2019; Brittain et al.,,
2019) is one of the most frequently used strategies to sample important
transitions. Its effectiveness in $Q$-learning is verified on Atari
environments. Recently, learning-based sampling, which utilizes neural
networks to generate priority scores, have shown to be outperform rule-based
PER (Zha et al.,, 2019; Sinha et al.,, 2020). However, ERO (Zha et al.,,
2019)’s gain over PER is not very significant, and (Sinha et al.,, 2020)
requires to determine some hyper-parameters and its effective on discrete
control tasks, e.g., Atari games, is unclear.

##  Appendix E Example of Implementating Our Method

To verify that our method is easy to implement, as we mentioned in the
manuscript, we explain how we can apply our method to Soft Actor Critic (SAC)
(Haarnoja et al.,, 2018b, a) based on PyTorch 1.7.1. In SAC, it uses double
critic networks and their target networks. We can implement model-augmented
Q-networks (MQNs) from the given critic networks as

---
Figure E.1: Implementation of MQNs. The parts with red lines are the newly
added parts, respectively.

In the update step of SAC, we can compute model errors as follows:

---
Figure E.2: Implementation of model errors in MQNs.

Then we can formulate MQN’s loss by summing model- and TD-errors as in Figure
E.3. The coefficient of each loss adaptively changes by a dynamic method in
(Liang & Zhang,, 2020). Furthermore, we use the quantity to impose priorities
to corresponding samples as in Figure E.4. These modifications are all for
applying model-augmented $Q$-learning (MQL). Note that there is no hyper-
parameters change. In a similar manner, we can apply MQL to other algorithms,
e.g., TD3 (Fujimoto et al.,, 2018), and Rainbow (Hessel et al.,, 2018; van
Hasselt et al.,, 2019). In the case of Rainbow, we use the context obtained by
convolutional layers.

---
Figure E.3: Formulating MQN’s loss.
---
Figure E.4: Formulating priorities of samples.

##  Appendix F Additional Experimental Results

We add efficient-Rainbow (van Hasselt et al.,, 2019)’s results on other Atari
games. MRainbow-MPER also outperforms baselines in all tested cases.

|  |  |
---|---|---|---
(a) Alien | (b) Asterix | (c) Gopher | (d) Kangaroo
Figure F.1: Learning curves of Rainbow on Atari games. Here, MRainbow-MPER is
a variant of Rainbow, where all of MQN, MPER, and MReward are applied.
Although Rainbow basically adopts PER, we denote it as Rainbow-PER for
consistency. One can observe that MRainbow-MPER overwhelms other methods. The
solid line and shaded regions represent the mean and standard deviation,
respectively, across five runs with random seeds.
---
Hopper-v3
Figure F.2: Learning curves of MBPO on Hopper-v3 belonging to MuJoCo
environments. One can observe that learning terminal functions is harmful to
sample efficiency in MBPO. The solid line and shaded regions represent the
mean and standard deviation, respectively, across five runs with random seeds.

Moreover, as we mentioned in Section D, to show that prior knowledge is
crucial to MBRL, we provide an experimental result about one of the state-of-
the-art MBRL methods, i.e., MBPO (Janner et al.,, 2019) with and without
learning a termination function. To verify it, we increased the output
dimension of an ensemble of dynamic models by one. Then we used the sigmoid
function to learn the probability of the termination function with the binary
cross-entropy loss. One can observe that the learning termination function is
significantly harmful to the sample efficiency compared to the original MBPO
in Figure F.2. The degraded result comes from wrongly computed TD-errors since
computing the TD target depends heavily on the termination function.