# THE BACKGROUND FIELD METHOD APPLIED TO COSMOLOGICAL PHASE TRANSITION

###### Abstract

In this paper, the cosmological phase transition is investigated by background
gauge field method. As a continuation of previous our work, some numerical
results and graphic solutions at $T\neq 0$ are presented. Hence the mechanism
of cosmological phase transition in the early Universe is considered. It is
shown that the breaking of symmetry significantly depend on the nonzero
temperature and chemical potential. Furthermore, it is the first order of
phase transition. Non - restoration of symmetry in hot gauge theories for
Cosmology.

PHAN HONG LIEN, NGUYEN NHU XUAN

Le Quy Don Technical University

##  1 INTRODUCTION

Phase transition is a complicate physical process, its nature is non
pertubative phenomenon [1] - [5]. Therefore, it is worth to mention that the
finite temperature effective action basing on functional integral as a method
which provides a general approximation beyond one loop and higher free energy
density in the pertubative as well as non-pertubative sector. In particular,
it plays an important role in the investigation of phase transition and non-
equilibrium phenomena [6],[7], as dynamics of processes.
Recently there has been considerable interest in the symmetry in both hot
scalar field theories and hot gauge theories for cosmology [8],[10]. It is
shown that the phase transition in the early Universe could be described by
non-Abelian gauge theories at high temperature [11]. Furthermore, the
contribution to the free energy only by calculated by using non-pertubative
methods. However, the effective potential of gauge theories may fail to be
gauge because it generally does depend on the $\xi$ \- gauge. Therefore the
background field method have been applied to compute quantum effects without
losing gauge invariance [12],[13].
Our main aim is to apply the background gauge field method at high temperature
to investigation of cosmological phase transition. In this connection, it is
possible to consider this paper as being complementary numerical computation
to result of previous our work.
This paper is organized as follows. In section II, we present the formalism of
effective action and background gauge field method. The one loop free energy
density at nonzero temperature an chemical potential is obtained in section
III. Section VI is devoted to investigation the cosmological phase transition
from some numerical calculations and graphic solutions. Our conclusion is
summarized in section V.

##  2 FORMALISM

We start from the Lagrangian density

$\begin{split}L_{0}=&-\frac{1}{4}F_{\mu\nu}^{a}F^{a\mu\nu}+\mathbf{\bar{\Psi}}(i\gamma^{\mu}D_{\mu}-G_{i}\Phi_{i})\mathbf{\Psi}\\\
&+\left[(D_{\mu}-i\mu\delta_{\mu
0})\Phi_{i}\right]^{+}\left[(D^{\mu}-i\mu\delta^{\mu
0})\Phi_{i}\right]-m^{2}\Phi^{+}_{i}\Phi_{i}-\lambda(\Phi^{+}_{i}\Phi_{i})^{2}\\\
&-\frac{1}{2\xi}\left(\partial^{\mu}A_{\mu}^{a}\right)^{2}-\partial_{\mu}\omega_{a}^{*}\partial^{\mu}\omega_{a}+f_{abc}(\partial_{\mu}\omega_{a}^{*})A_{\mu}^{b}\omega^{c}.\end{split}$
(1)

where $\mathbf{\bar{\Psi}},\mathbf{\Psi}$ are multiplet of fermion fields,
$\Phi_{i}(i=1,2\ldots n)$ are components of scalar fields, $A_{\mu}$ \- gauge
fields and $\omega,\omega^{*}$ \- ghost fields. Here $\mu$ is chemical
potential, $G_{i}$ and $\lambda$ are coupling constants, $\lambda>0$

$\begin{split}D_{\mu}\equiv&\hskip 3.0pt\partial_{\mu}-iT^{a}A_{\mu}^{a},\\\
F_{\mu\nu}^{a}\equiv&\hskip
3.0pt\partial_{\mu}A_{\nu}^{a}-\partial_{\nu}A_{\mu}^{a}+f_{abc}A_{\mu}^{b}A_{\nu}^{c},\end{split}$

where $T_{a}$ are group generators, $f_{abc}$ are structure constants which
satisfy Lie algebra

$\displaystyle f_{abc}f_{dbc}=$ $\displaystyle\hskip
3.0ptg^{2}C_{A}\delta_{ad}$ (2) $\displaystyle Tr(T_{a}T_{b})=$
$\displaystyle\hskip 3.0ptg^{2}C_{F}\delta_{ab}$ (3)

with $C_{A}$ is numerical constant of gauge group, $C_{A}=N$ for $SU(N)$,
$C_{F}$ is representation of this group.
The fields are shifted by

$\displaystyle A_{\mu}\rightarrow$ $\displaystyle\hskip
3.0pt\mathbf{A}_{\mu}+A^{\prime}_{\mu};\hskip 28.45274pt\langle
0|\mathbf{A}_{\mu}|0\rangle=const,\hskip 14.0pt\langle
0|A^{\prime}_{\mu}|0\rangle=0,$ (4) $\displaystyle\Phi_{i}\rightarrow$
$\displaystyle\hskip 3.0pt\mathbf{\Phi}_{i}+\Phi^{\prime}_{i};\hskip
28.45274pt\hskip 5.0pt\langle 0|\mathbf{\Phi}_{i}|0\rangle\hskip
3.0pt=\phi_{0i},\hskip 28.0pt\langle 0|\Phi^{\prime}_{i}|0\rangle=0,$ (5)
$\displaystyle\Psi\rightarrow$ $\displaystyle\hskip
3.0pt\mathbf{\Psi}+\Psi^{\prime};\hskip 28.45274pt\hskip 9.0pt\langle
0|\mathbf{\Psi}|0\rangle\hskip 4.0pt=\langle
0|\mathbf{\Psi}^{\prime}|0\rangle=0,$ (6) $\displaystyle\omega_{a}\rightarrow$
$\displaystyle\hskip 3.0pt\boldsymbol{\omega}_{a}+\omega^{\prime}_{a};\hskip
28.45274pt\hskip 6.0pt\langle 0|\boldsymbol{\omega}_{a}|0\rangle=\langle
0|\omega^{\prime}_{a}|0\rangle=0,$ (7)

where $\mathbf{A}_{\mu},\mathbf{\Phi},\mathbf{\Psi},\boldsymbol{\omega}_{a}$
are the background fields, and
$A^{\prime}_{\mu},\Phi^{\prime},\Psi^{\prime},\omega^{\prime}_{a}$ are the
quantum fields, which are variables of integration in the functional integral.
It is well known, the background field method allows one to fix a gauge,
thereby compute quantum effects without losing explicit gauge invariance [14].
Now we consider the effective action in background field for which
$\mathbf{A}_{\mu}^{a}$ and $\mathbf{\Phi}_{i}$ are constant, and
$\mathbf{\Psi}=\mathbf{\bar{\Psi}}=\boldsymbol{\omega}=\boldsymbol{\omega}^{*}=0$.
The effective action is calculated from the part of the action that is
quadratic in quantum fields $A^{\prime}_{\mu},\Phi^{\prime},\Psi^{\prime}$ and
$\omega^{\prime},\omega^{\prime*}$ over which one integrated

$\begin{split}I_{quad}=&\int dxL_{quad}=\int
dx\left[-\frac{1}{4}\left(\bar{D}_{\mu}A_{\nu}^{{}^{\prime}a}-\bar{D}_{\nu}A_{\mu}^{{}^{\prime}a}\right)^{2}-\frac{1}{4}F_{\mu\nu}^{a}f_{abc}A_{\mu}^{{}^{\prime}b}A_{\nu}^{{}^{\prime}c}\right]\\\
&-\int
dx\bar{\Psi^{\prime}}\left(\gamma^{\mu}D_{\mu}+\mu\gamma^{0}+M+g_{i}\phi^{\prime}_{i}\right)\Psi^{\prime}\\\
&+\int
dx\left[\frac{1}{2}\left(D_{\mu}\phi^{\prime}_{i}D^{\mu}\phi^{\prime}_{i}-M_{ij}^{2}\phi_{i}^{{}^{\prime}2}\right)-\frac{\lambda}{4}\phi_{i}^{{}^{\prime}4}\right]\\\
&-\int
dx\left[\frac{1}{2\xi}\left(\bar{D}_{\mu}A_{\nu}^{{}^{\prime}a}\right)^{2}+\left(\bar{D}_{\mu}\omega_{a}^{{}^{\prime}*}\right)\left(\bar{D}_{\mu}\omega^{\prime}_{a}\right)\right]\end{split}$
$\begin{split}I_{quad}=&\frac{1}{2}\int
dxdyA_{\mu}^{{}^{\prime}a}(x)\mathfrak{D}_{\mu\nu}^{ab}(x,y)A_{\nu}^{{}^{\prime}b}(y)-\int
dxdy\Psi^{\prime}(x)\mathfrak{D}_{ab}(x,y)\Psi^{\prime}(y)\\\
&+\frac{1}{2}\int
dxdy\phi^{\prime}_{i}(x)\mathfrak{D}_{ik}(x,y)\phi^{\prime}_{k}(y)-\int
dxdy\omega_{a}^{{}^{\prime}*}(x)\mathfrak{D}_{ab}(x,y)\omega^{\prime}_{b}(y)\end{split}$
(8)

By using the Fourier transformation $\mathfrak{D}(k)=\int
dxe^{ik(x-y)}\mathfrak{D}(x-y)$ the matrices in (8) are given by

$\begin{split}\mathfrak{D}_{\mu\nu}^{ab}(k)=&g_{\mu\nu}\Bigr{[}\left(-ik_{\rho}\delta_{ca}+f_{cda}A_{\rho}^{d}\right)\left(-ik^{\rho}\delta_{cb}+f_{ceb}A_{\rho}^{e}\right)\\\
&-\left(-ik_{\nu}\delta_{ca}+f_{cda}A_{\nu}^{d}\right)\left(-ik_{\mu}\delta_{cb}+f_{ceb}A_{\mu}^{e}\right)+F_{\mu\nu}^{c}f_{cab}\Bigr{]}\\\
&+g_{\mu\nu}\delta_{ij}\Phi_{i}(k)\Phi_{j}(k)T^{a}T^{b}+\mbox{$\epsilon$
terms},\end{split}$ (9)

with $F_{\mu\nu}^{a}=f_{abc}A_{\mu}^{b}A_{\nu}^{c}$.

$\displaystyle\mathfrak{D}(k)=$ $\displaystyle(-ik/-iT^{a}A^{a}_{\mu}/\hskip
9.0pt+M+\mu\gamma_{0})+\mbox{$\epsilon$ terms},$ (10)
$\displaystyle\mathfrak{D}_{ab}(k)=$
$\displaystyle\left(-ik_{\rho}\delta_{ca}+f_{cda}A_{\rho}^{d}\right)\left(-ik^{\rho}\delta_{cb}+f_{ceb}A^{\rho
b}\right)+\mbox{$\epsilon$ terms},$ (11) $\displaystyle\mathfrak{D}_{ij}(k)=$
$\displaystyle\left(-ik_{\mu}-i\mu-
iT_{a}A_{\mu}^{a}\right)_{i}\left(ik_{\mu}+i\mu+iT_{a}A_{\mu}^{a}\right)_{j}-m_{ij}^{2}-\frac{\lambda}{2}\phi_{i}\phi_{j}+\mbox{$\epsilon$
terms},$ (12)

In momentum representation, the effective action takes the general form

$\begin{split}&\Gamma_{\beta}\left[\psi,\bar{\psi},\phi,A_{\mu},\omega,\omega^{*}\right]=I\left[\psi,\bar{\psi},\phi,A_{\mu},\omega,\omega^{*}\right]-\frac{i}{2}TrlnG_{\mu\nu}^{ab}(k)\\\
&+iTrlnS(k)-\frac{i}{2}Trln\Delta_{ij}(k)+iTrlnD_{ab}(k)+\sum_{n=2}^{\infty}\mbox{\emph{n
loops 1PI}}.\end{split}$ (13)

Here the action $I\left[\psi,\bar{\psi},\phi,A_{\mu},\omega,\omega^{*}\right]$
is given in (8) - (12). The Trace, the logarithm are taken in functional
sense, and the free propagators are given by

$\displaystyle S^{-1}(k)$ $\displaystyle=k/-M-i\epsilon;\hskip
71.13188ptM=g\nu,$ (14) $\displaystyle\Delta_{ij}^{-1}(k)$
$\displaystyle=\delta_{ij}k^{2}-M_{ij}^{2}-i\epsilon;\hskip
45.52458ptM_{ij}^{2}=(\mu^{2}-m^{2})\delta_{ij}-\frac{\lambda}{2}\phi_{i}\phi_{j},$
(15) $\displaystyle\left[G_{0ab}(k)\right]_{\mu\nu}^{-1}$
$\displaystyle=\left(M^{2}_{ab}-k^{2}\delta_{ab}\right)\left[\frac{k_{\mu}k_{\nu}}{k^{2}}-g_{\mu\nu}\right]+\left[\delta_{ab}\frac{k^{2}}{\xi}-M^{2}_{ab}\right]\frac{k_{\mu}k_{\nu}}{k^{2}},$
(16) $\displaystyle D_{0ab}^{-1}(k)$
$\displaystyle=\delta_{ab}(k^{2}-i\epsilon),\hskip
71.13188ptM_{ab}=\frac{1}{2}\delta_{ab}g\nu.$ (17)

##  3 ONE LOOP THERMAL FREE ENERGY DENSITY

Next we consider the theory at finite temperature by ”imagine time” formalism.
For
$\langle\mathbf{A}_{\mu}\rangle_{\beta}=\delta_{0\mu}A_{\mu}^{0},\langle\mathbf{\Phi}\rangle=\phi_{0}$
the effective potential is defined by

$V_{\beta}=-\frac{\Gamma_{\beta}}{\beta\int dx}$ (18)

where $\beta=T^{-1}$ (we set Boltzmann constant k = 1). It is just thermal
free energy density, which concerns with the phase transition at $T=T_{c}$.
In $d=4-2\epsilon$ dimension, the divergent integrals can be regularized by
using the ”imagine time” formalism, where nonzero chemical potential $\mu$ is
added to the fermionic Matsubara frequencies, i.e.

$i\int\frac{d^{4}k}{(2\pi)^{4}}f(k)\rightarrow
T\sum\int\frac{d^{3}k}{(2\pi)^{3}}f(i\omega_{n},\vec{k}),$

where $i\omega_{n}=2\pi nT$ for bosons, $\omega_{n}=(2n+1)\pi T$ for fermions
and $i\omega_{n}\rightarrow i\omega_{n}+\mu$.
From (8) - (13) and (18) we arrived at the expression for the effective
potential

$\begin{split}V_{\beta}=&V_{cl}-\sum_{k}\int\left[ln(k^{2}+M^{2})-ln(k^{2}+M_{ij}^{2})-ln(k^{2}+M_{ab}^{2})\right]\\\
&+ig^{2}\left(\frac{11}{12}N-\frac{1}{6}N_{F}+\frac{1}{12}N_{B}\right)\sum_{k}\int\frac{1}{(k^{2})^{2}}\int
dxF_{\mu\nu}^{a}F^{a\mu\nu}\\\
&+\frac{i}{8}g^{2}\sum_{k}\int\frac{1}{(k^{2}+m^{2})^{2}}\\\
&+\frac{i}{2}g_{i}\sum_{k}\int\sum_{p}\int\frac{1}{(p^{2}+M_{1}^{2})(k^{2}+M_{2}^{2})\left[(k+p)^{2}+M_{ab}^{2}\right]}.\end{split}$
(19)

Finally, the one loop thermal free energy density is given

$\displaystyle V_{\beta}=$
$\displaystyle-\frac{1}{2}(\mu^{2}-m^{2})\phi^{2}+\frac{\lambda}{4}\phi^{4}-\frac{1}{2}\delta_{ab}M_{ab}^{2}A_{0\mu}^{2}-\frac{\pi^{2}T^{4}}{90}\left(N_{B}+\frac{7}{8}N_{F}\right)$
$\displaystyle+\frac{T^{2}}{24}\left\\{(\mu^{2}-m^{2}-\frac{\lambda}{2}\phi^{2})+3TrM_{ab}^{2}+\frac{1}{2}Tr\left[\gamma_{0}(M+\mu\gamma_{0})\gamma_{0}(M+\mu\gamma_{0})\right]\right\\}$
$\displaystyle-\frac{T}{12\pi}\left(\mathfrak{M}^{3}+\delta_{ab}M_{ab}^{3}\right)-\frac{g^{2}T^{3}}{48\times
4\pi}\left(\mathfrak{M}+\delta_{ab}M_{ab}+2M\right)$ (20)
$\displaystyle+\frac{g^{2}}{(4\pi)^{2}}\left(\frac{11}{12}N-\frac{1}{6}N_{F}+\frac{1}{12}N_{B}\right)\left(\frac{1}{\epsilon}-2ln\frac{\bar{\eta}}{4\pi
T}+2\gamma_{E}\right)\int dxF_{\mu\nu}^{a}F^{a\mu\nu}.$

where $\mathfrak{M},M_{ab},M$ are thermal masses, e.g the squared scalar mass
is

$\mathfrak{M}^{2}=(\mu^{2}-m^{2})+\frac{\lambda}{24}T^{2}-\frac{\lambda}{2}\phi^{2}.$
(21)

The renormalized coupling is given by

$g_{R}=g\left[1+\frac{g^{2}}{4\pi^{2}}\left(\frac{11}{12}N-\frac{1}{6}N_{F}+\frac{1}{12}N_{B}\right)\left(\frac{1}{\epsilon}+2ln\frac{\bar{\eta}}{4\pi
T}+2\gamma_{E}\right)\right]+0(g^{4})$ (22)

i.e the physical coupling $g_{R}$ increases due to quantum corrections in the
non - Abelian theory.

##  4 SOME NUMERICAL RESULTS AT NONZERO TEMPERATURE AND CHEMICAL POTENTIAL

Let us consider the minimum of the effective potential from which we can
determine the breaking or the restoration of symmetry.
In the case of $\mu\leq m$, the free energy density is minimum at $\Phi_{0}=0$
at any temperature T, i.e the symmetry is not broken (Fig. 1).
In Fig. 2 it is easily seen that the high temperature and nonzero chemical
potential significantly affect on the spontaneous breaking of symmetry. The
importance is that it leads to the cosmological phase transition.
In Fig. 3, we can see that the effective squared scalar mass is change from
negative to positive value when the temperature is high enough. It is due to
the contribution of the $\mathfrak{M}^{2}T^{2}$ term in the effective
potential. This just is phenomena of breaking of symmetry, when it occurs the
cosmological phase transition is manifested.
In Fig. 4, the free energy is represented as a continuous function of $T$ and
$\mu$ when $\mu\leq m$, but in Fig. 5, it has a discontinuity when $\mu\geq
m$. It is shown by the discontinuous buffer between the symmetry and its
breaking part. This is just the cosmological phase transition. Furthermore,
the restoration of symmetry does appear that means after symmetry was
spontaneously broken, the Universe is asymmetry
In Fig.6, the effective squared scalar mass depend on the chemical potential
$\mu$ at fixed temperature, which expresses the first order phase transition.

|
---|---
a) |  b)

Fig.1. The effective potential $V_{\beta}$ as a function of T and $\Phi$ in
the case $\mu/m=0.5;1.0$.
a) $\Phi=-100\div 100MeV$ with $T=0,100,200,400MeV$ respectively
b) $T=0\div 400$ MeV and $\Phi=-300\div 300$ MeV

|
---|---
a) |  b)

Fig.2. The effective potential $V_{\beta}$ as a function of T and $\Phi$ in
the case $\mu/m=1.5$.
a) $\Phi=-100\div 100MeV$ with $T=0,100,200,400MeV$, respectively
b) $T=0\div 400$ MeV and $\Phi=-300\div 300$ MeV.

(a)

|
---|---
b) | c)

Fig.3. The squared scalar mass $\mathfrak{M}^{2}$ as a function of temperature
T.

a) $T=-400\div 400MeV$ with $\Phi=100MeV$ in cases $\mu=m,\mu=1,5m,\mu=0,5m$.
b,c) $T=-400\div 400$ MeV and $\Phi=-300\div 300$ MeV in cases
$\mu=1,2m,\mu=0,7m$, respectively.

|
---|---
a) |  b)

Fig.4. The effective potential $V_{\beta}$ as a function of $\mu$ and $\Phi$
at $T=400MeV$.
a) $\mu=-500\div 500MeV$ with $\Phi=0,100,200,400MeV$, respectively.
b) $\mu=0\div 400$ MeV and $\Phi=-300\div 300$ MeV.

Fig.5. The effective mass as a function of chemical potential. It expresses
the first order
phase transition.The discontinuity is phase transition region (white region).
Plot for $\mu=0\div 500$ MeV and $\Phi=-300\div 300$ MeV.

Fig.6. Dependence of the effective squared scalar mass on $\mu$ at fixed
temperature.
It expresses the first order phase transition.
Plot for $\mathfrak{M}^{2}$ with $\mu=0\div 1200$ MeV.

##  5 DISCUSSION AND CONCLUSION

We have studied the free energy as a function of temperature and non-zero
chemical potential by background gauge field method in frame of Abelian
theories. Hence the mechanism of cosmological phase transition is
investigated. The graphic solutions have shown that in the early Universe it
is first order phase transition. Furthermore, after the spontaneous symmetric
breaking, one also can see non-restoration of symmetry in hot gauge field
theories for Cosmology, i.e the Universe is asymmetry.
We should mention that one can study the quantum effects in phase transition
and find exactly the cosmological critical temperature if a suitable parameter
set is given.

## ACKNOWLEDGMENT

One of the authors (PHL) would like to thank Prof. Tran Huu Phat for helpful
suggestion of this problem.

## References

  * [1] L. Doland and R. Jackiw, Phys. Rev. D9, 3321 (1974).
  * [2] S.Weinberg, Phys. Rev. D9, 3357 (1974).
  * [3] J. I. Kapusta, _Finite Temperature Field Theory_ , Cambridge University Press, Cambridge (1989).
  * [4] N. Goldenfeld. _Lectures on Phase Transition and Renormalization Group_ , Addison Wesley, NewYork (1972).
  * [5] M. Le Bellac, _Thermal Field Theory_ , Cambridge University Press, Cambridge (1996).
  * [6] H. Umezawa, _Equilibrium and non equilibrium thermal physics_ , in _Banff/Cap Workshop on Thermal Field Theory_ , eds. F.C. Khanna, R. Kobes, G.Kunstatter and H.Umezawa. World Scientific Singapore (1994).
  * [7] T. Kunihiro, T. Muto, T. Takasuka, R. Tamagaki and T. Tatsumi, _Various Phases in High- Density Nuclear Matter and Neutron Stars_ , Prog. Theor. Phys. Supplement 112 (1993).
  * [8] M.P. Lombardo, _Finite Temperature Field Theory and Phase Transitions_ , ICTP Lecture notes on Astroparticle Physics and Cosmology (2001).
  * [9] Marcus B. Pinto, Rudnei O. Ramos and Frederio F. de Souza Cruz. Nucl. Phys. B (Proc. Suppl.) 121 (2003).
  * [10] E. W. Kolb and M. S. Turner, _The Early Universe_ , Addison Wesley (1990).
  * [11] P.Arnold, O. Espinosa, Phys. Rev, D17, (1994) 3546
  * [12] B. S. De Witt, _A gauge invariant effective action_ , in Quantum Gravity II, Oxford: Oxford University Press (1981).
  * [13] R. Jackiw, G. Amelia - Camelia, _Field theoretical background for thermal physics_ , Third workshop on Thermal Field Theories, Banff, Canada (1993).
  * [14] S. Pokorski, _Gauge Field Theory_ , Cambridge University Press, Cambridge (1987).
  * [15] P. Arnold, C. Zhai, Phys. Rev. D49, 2740 (1974).


# Sequence Modeling of Temporal Credit Assignment for Episodic Reinforcement
Learning

Yang Liu
UIUC
&Yunan Luo
UIUC
&Yuanyi Zhong
UIUC
&Xi Chen
covariant.ai
&Qiang Liu
UT Austin
&Jian Peng
UIUC

###### Abstract

Recent advances in deep reinforcement learning algorithms have shown great
potential and success for solving many challenging real-world problems,
including Go game and robotic applications. Usually, these algorithms need a
carefully designed reward function to guide training in each time step.
However, in real world, it is non-trivial to design such a reward function,
and the only signal available is usually obtained at the end of a trajectory,
also known as the episodic reward or return. In this work, we introduce a new
algorithm for temporal credit assignment, which learns to decompose the
episodic return back to each time-step in the trajectory using deep neural
networks. With this learned reward signal, the learning efficiency can be
substantially improved for episodic reinforcement learning. In particular, we
find that expressive language models such as the Transformer can be adopted
for learning the importance and the dependency of states in the trajectory,
therefore providing high-quality and interpretable learned reward signals. We
have performed extensive experiments on a set of MuJoCo continuous locomotive
control tasks with only episodic returns and demonstrated the effectiveness of
our algorithm.

##  1 Introduction

Deep reinforcement learning (RL) methods, including the well-known policy
gradient algorithms [24, 31, 32] and deep Q-networks [25], have shown superior
performance and great potential in many difficult real-world problems, such as
the Go game [33, 34], locomotive continuous control problems [21], resource
management [23], and robotics [20]. The key idea of such algorithms is to use
deep neural networks as functional approximators to abstract or represent
complex state observation so that actions can be properly chosen accordingly
to optimize a long-term expected return. The learned policy or Q function
essentially captures the temporal structure of the sequential decision problem
and decompose it to a supervised learning problem, guided by the reward
signal. However, in many real-world problems, the reward signal is usually not
dense enough to provide sufficient supervision for learning the decision at
each single time step. In many practical tasks, such as the Go game and the
automatic chemical design problems [27], we can only obtain a final reward or
return value after finishing the entire rollout of the policy, while no
intermediate reward is provided before reaching the end of the trajectory.
This type of problems is also known as the episodic reinforcement learning.

Unfortunately, when the reward signal becomes delayed or even episodic, most
existing deep reinforcement learning algorithms may get stuck during the
training process and often suffer from inferior performance and inefficient
sample complexity [8, 9]. This problem is widely known as the temporal credit
assignment in reinforcement learning [39], which describes the issue of
delayed rewards causing the signal to be diluted over time and only weakly
affecting the states temporally distant from the time step when the rewards
get collected. For example, in Go games, the only effective reward is the
final win or loss. This reward is received only after finishing the entire
games, usually consisting of hundreds of moves in the trajectory. During the
game, though human professional players or experts may be able to decide which
moves are likely to influence the final winning probability, quantitatively
design of such informative rewards is very challenging. Such sparse or
episodic reward signal makes the training of policies or Q-neural networks
very difficult, as a lot of data is required to propagate the final win/loss
reward back to earlier states or state-action pairs. In addition, since there
is no immediate reward or even no short-term reward, the exploration becomes
challenging without any information. Similar to Go games, there are many real-
world problems in which rewards are terribly delayed or episodic, with only a
non-zero reward value received at the end of the episode or the trajectory.

In this paper, we introduce a new algorithm for the temporal credit assignment
problem. The idea is to learn deep neural networks that is able to decompose
the episodic reward into parts and assign them back to each time step in the
trajectory. After learning this decomposition, we can use the assigned dense
reward signal to guide policy optimization. Formally, we derive a generalized
policy gradient for decomposed reward signal. In our derivation, we find that
to ensure the correctness of the policy gradient algorithms, the dependency of
the neural network on time steps needs to be forward-looking, i.e. for a time
$t$ before the terminal step $T$, the reward function
$r_{t}(\\{s_{t^{\prime}},{a_{t}^{\prime}}\\}_{t^{\prime}=0}^{T})=r_{t}(\\{s_{t^{\prime}},a_{t^{\prime}}\\}_{t^{\prime}=0}^{t})$
does not depend on any future states or actions. This forward structure,
frequently seen in natural language processing [18], motivates us to apply the
renowned neural-network sequence models, such as language models to learning
the reward decomposition. In particular, we adopt the Transformer, which
satisfies the forward-looking structure, for learning the importance and the
dependency of states using the self-attention mechanism. With the learned
reward signal, we show on a set of MuJoCo continuous locomotive control tasks
with episodic returns that the learning and sample efficiency of can be
greatly improved.

##  2 Method

###  2.1 Background and Problem Statement

#### Reinforcement learning

Reinforcement learning considers the problem of finding an optimal policy for
an agent which interacts with an environment and collects reward per action.
The goal of the agent is to maximize the cumulative reward along a trajectory.
Formally, this problem can be formulated as a Markov decision process (MDP)
over the environment states $s\in{\it S}$ and agent actions $a\in{\it A}$,
under an unknown environmental dynamics defined by a transition probability
$T(s^{\prime}|s,a)$. The agent’s action $a$ is selected by a conditional
probability distribution $\pi_{\theta}(a|s)$ parameterized by
$\theta\in\Theta$. Playing the policy repeatedly under MDP yields a trajectory
$\tau=\\{s_{t},a_{t}\\}_{t=0}^{T}$, where $T$ denotes the horizon length. Each
trajectory $\tau$ is associated with a reward function $R(\tau)$ which we want
to optimize, that is,
$\max_{\theta}J(\theta):=\mathbb{E}_{\pi_{\theta}}\left[R(\tau)\right],$ where
the expectation $\mathbb{E}_{\pi_{\theta}}$ is with the distribution of the
unknown dynamics $T(s_{t+1}|s_{t},a_{t})$ and policy
$\pi_{\theta}(a_{t}|s_{t})$.

In standard RL settings, it is common to assume that the reward is a
discounted sum of a set of local reward functions distributed across the time,
that is, $R(\tau)=\sum_{t=0}^{\infty}\gamma^{t}r_{t}\left(s_{t},a_{t}\right),$
where the local reward signal $r(s,a)$ is assumed to be observed immediately
following the action $a$ performed at state $s$, and $\gamma\in[0,1)$ is the
discount factor. This decomposition structure greatly simplifies the problem,
and forms the basic assumption of popular policy gradient algorithms. The
focus of this work, however, is the more difficult case when $R(\tau)$ does
not have a simple or known decomposition structure priori, and is observed
only after the whole trajectory $\tau$ is rolled out.

Policy Gradient There are several different types of algorithms for learning
policies, including Q-learning, policy gradient and evolutionary algorithms.
For the general episodic reward function $R(\tau)$, we have a basic gradient
estimation derived using the likelihood ratio trick,

$\displaystyle\begin{split}\nabla_{\theta}J(\theta)&=\nabla_{\theta}\mathbb{E}_{\pi_{\theta}}[R(\tau)]=\mathbb{E}_{\pi_{\theta}}\left[R(\tau)\sum_{t=0}^{T}\nabla_{\theta}\log\pi(a_{t}|s_{t})\right].\end{split}$
(1)

This basic algorithm, however, often yields large variance in gradient
estimation. The policy gradient theorem [38] allows us to derive a simplified
formula for the case when the reward is decomposed step-wisely:

$\nabla_{\theta}J(\theta)=\mathbb{E}_{\pi}\left[\nabla_{\theta}\log\pi(a|s)Q^{\pi}(s,a)\right],$
(2)

where
$Q^{\pi}(s,a)=\mathbb{E}_{\pi}\left[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})|s_{0}=s,a_{0}=a\right]$
denotes the expected return under policy $\pi$ starting from state $s$ and
action $a$. With empirical estimate of $\hat{Q}^{\pi}(s_{t},a_{t})=\sum_{j\geq
t}\gamma^{j-t}r_{j}$ using rollout trajectories, then we can obtain the well-
known REINFORCE policy gradient [41] as

$\hat{\nabla}_{\theta}J(\theta)=\frac{1}{T}\sum_{t=0}^{T}\gamma^{t}\nabla_{\theta}\log\pi(a_{t}|s_{t})\hat{Q}^{\pi}(s_{t},a_{t}).$
(3)

Improved policy gradient methods, such as the Proximal Policy Optimization
(PPO) [32, 12], are now able to provide the state-the-art performance on many
problems. It uses a proximal Kullback-Leibler (KL) divergence penalty to
regularize and stabilize the policy gradient. Furthermore, control variate
methods, such as the generalized advantage estimation (GAE) [31], help
reducing the variance of the policy gradient estimation. So far, most policy
gradient methods are designated for infinite-horizon, dense reward settings,
as the dense rewards can provide direct supervision for value function
estimation and policy improvement in each time step.

Episodic RL For tasks with episodic rewards, usually, there is a set of
terminal states. At the end of each trajectory $\tau$, a reward is only
received at the terminal state. In other words, before reaching the final
state $s_{T}$, rewards $r_{t}(s_{t},a_{t})=0$ for all $t<T$. In many tasks,
the terminal states can be predefined, or the length of the trajectories is
limited. For simplicity, we omit the discount factor and assume the trajectory
length is at most $T$. In this way, we can abuse the notation of $s_{T}$ to
denote the last state without further confusion. Therefore, the objective of
episodic reinforcement learning becomes
$J(\theta)=\mathbb{E}_{\pi_{\theta}}\left[R(\tau)\right]=\mathbb{E}_{\pi_{\theta}}\left[r_{T}(s_{T})\right].$
Note that it is not hard to add the discount factor back and/or adopt the
mathematical definitions and derivations for problems with a set of terminal
states.

For episodic problems, the straightforward application of policy gradient
methods, including REINFORCE [41], A2C [24] and PPO [32], may suffer from
sample inefficiency, as the final episodic reward would only provide the same
or similar supervision for learning policy over all time steps in a
trajectory. Therefore, a huge volume of rollout trajectories are required to
distinguish the subtle influence of certain action on the final reward.
Besides policy optimization methods, blackbox optimization approaches , such
as cross entropy method [29], CMA-ES [11] and evolution strategies [30], have
been also applied to episodic RL problems due to their computational
efficiency.

###  2.2 Overview of Our Approach

We propose a new approach to learn a dense surrogate reward function that
approximates the temporal credit assignment of the episodic reward (Algorithm
1). The idea is intuitive: we hope to find $\hat{r}(s_{t},a_{t})$
approximating $R(\tau)=\sum_{t=0}^{T}\hat{r}(s_{t},a_{t})$, so that we can use
$\hat{r}$ as the surrogate reward to compute the policy gradient. If $\hat{r}$
is dense over time step and provide sufficient information about the influence
on the episodic reward, then it should help improve the sample efficiency of
training policies under the episodic reward setting.

Algorithm 1 Policy optimization with decomposed reward

1: Initialize: policy parameters $\theta_{0}$, predictor parameters $\phi_{0}$

2: for $i$ = 1, 2, 3, … $N$ do

3: Collect a batch of trajectories using roll-outs.

4: Append the new trajectories into trajectory buffer for regression.

5: Train reward predictor using gradient descent:
$\phi_{i}\leftarrow\phi_{i-1}-\gamma_{\phi}\nabla_{\phi}{\mathcal{L}_{regression}}$

6: Update policy parameters using policy optimization algorithm:
$\theta_{i}\leftarrow\theta_{i-1}+\gamma_{\theta}\nabla J(\theta)$ where
$\nabla J(\theta)$ is obtained with Eq. (7) .

7: end for

8: Output: policy $\pi_{\theta_{N}}$, reward predictor $\phi_{N}$

Here, we consider a generalization of the time step-wise reward function.
Assume a reward function $\hat{r}$ that is defined on states and actions over
a time interval $\alpha\in\mathcal{I}$. Then we expect that the episodic
reward $R(\tau)=r_{T}(s_{T})$ can be decomposed as the sum of the reward
function on all intervals:
$\sum_{\alpha\in\mathcal{I}}\hat{r}(s_{\alpha},a_{\alpha})$, where
$s_{\alpha}=\\{s_{i}|i\in\alpha\\}$ and $a_{\alpha}=\\{a_{i}|i\in\alpha\\}$.
The choice of $\mathcal{I}$ can be very flexible: If each interval only
contains a single time step, then the reward function is defined on each time
step $\alpha\in\\{\\{0\\},\\{1\\},\ldots,\\{T\\}\\}$; If $\mathcal{I}$
contains all consecutive sub-sequences starting from time $0$, then each
$\alpha\in\\{\\{0,1,2,\ldots,t\\}:\forall t=0,\ldots,T\\}$ contains all time
steps from the beginning of the trajectory to the current time step $t$.

Therefore, the objective function of learning such reward $\hat{r}$ can be
done by minimizing the regression loss as following:

$\min_{\phi}\mathcal{L}_{regression}(\phi):=\sum_{\tau\in
D}(\sum_{\alpha\in\mathcal{I}}\hat{r}_{\phi}(s_{\alpha},a_{\alpha})-R(\tau))^{2}$
(4)

where $\hat{r}_{\phi}$ can be a neural network that takes $s_{\alpha}$ and
$a_{\alpha}$ as input and is parameterized by $\phi$, $D$ is a collection of
trajectories. There are several critical choices: first, how one should define
the interval set $\mathcal{I}$; second, which model would be used for the
reward function; third, how to collect the dataset $D$. We will discuss the
design principles in the following sections and their impact in the
experimental section.

###  2.3 Generalized Policy Gradient with Rewards on Time Intervals

#### Policy Gradient of Composite Reward

Now, let us assume we have learned a composite approximation of the reward
function,
$\hat{R}(\tau)=\sum_{\alpha\in\mathcal{I}}\hat{r}(s_{\alpha},a_{\alpha})$,
where $\hat{r}(s_{\alpha},a_{\alpha})$ is a local reward function that defined
on state and actions over time interval $\alpha\in\mathcal{I}$. Our key idea
is to leverage the decomposition structure of $\hat{R}$ to simplify and reduce
the variance of the policy gradient formula, as we summarize in the following
generalization of policy gradient theorem to the composite rewards.

###### Theorem 1

I) Denote by $\hat{J}(\theta):=\mathbb{E}_{\pi_{\theta}}[\hat{R}(\tau)]$ the
expectation of the composite reward $\hat{R}(\tau)$, we have

$\displaystyle\nabla\hat{J}(\theta)$
$\displaystyle=\nabla_{\theta}\mathbb{E}_{\pi_{\theta}}[\hat{R}(\tau)]=\sum_{\alpha\in\mathcal{I}}\mathbb{E}_{\pi_{\theta}}\left[\hat{r}(s_{\alpha},a_{\alpha})\sum_{t\in\Gamma_{\alpha}}\nabla_{\theta}\log\pi(a_{t}|s_{t})\right],$
(5)

where $\Gamma_{\alpha}=\\{t\colon t\leq\max(\alpha)\\}$ and $\max(\alpha)$
denotes the maximum element of set $\alpha$; note that $\Gamma_{\alpha}$ is
the set of all $t$ that $\nabla_{\theta}\log\pi(a_{t}|s_{t})$ should multiply
by $\hat{r}(s_{\alpha},t_{\alpha})$.

II) Equivalently, we have

$\nabla\hat{J}(\theta)=\mathbb{E}_{\pi_{\theta}}\left[\sum_{t=0}^{T}Q_{t}(\tau)\nabla_{\theta}\log\pi(a_{t}|s_{t})\right],$
(6)

where $Q_{t}$ is a generalized $Q$-function, defined to be

$Q_{t}(\tau)=\sum_{\alpha\in\Gamma_{t}^{*}}\hat{r}(s_{\alpha},a_{\alpha}),~{}~{}~{}\Gamma_{t}^{*}=\\{\alpha\colon\max(\alpha)\geq
t\\}.$

Here $\Gamma_{t}^{*}$ is the set of all $\alpha$ whose upper bound
$\max(\alpha)$ excess $t$.

Theorem 1 allows us to leverage decomposition structure of $\hat{R}(\tau)$ to
increase the efficiency of policy gradient. Compared to the basic gradient
formula in (1), Eq. (5) keeps only the $\nabla\log\pi(a_{t}|s_{t})$ in set
$\Gamma_{\alpha}$ for each local reward $\hat{r}(s_{\alpha},a_{\alpha}).$ This
is obtained by the fact of MDP that present actions only impact the future but
not the past, similar to what we have seen in the original policy gradient.

Eq. (6) in Theorem 1 can be viewed as a generalization of policy gradient
theorem in Eq. 2, and $Q_{t}$ is similar to the definition of typical notion
of $Q$-function, which corresponds to our special case when each
$\alpha\in\mathcal{I}$ includes an individual time steps, that is,
$\mathcal{I}=\\{\\{0\\},\\{1\\},\ldots,\\{T\\}\\}$.

By replacing the expectation in Eq. (6) with empirical average from rollout
trajectories, we can derive a generalized policy gradient for the composite
reward $\hat{R}(\tau),$ that is, given a set of trajectories
$\tau^{i}=\\{(s_{t}^{i},a_{t}^{i})_{t=0}^{T}\\}$, $i=1,\ldots,n$, we have

$\nabla_{\theta}\hat{J}(\theta)=\frac{1}{n}\sum_{i=1}^{n}\sum_{t=0}^{T}Q_{t}(\tau^{i})\nabla\log\pi(a_{t}^{i}|s_{t}^{i}),$

where again
$Q_{t}(\tau^{i})=\sum_{\alpha\in\Gamma_{t}^{*}}\hat{r}(s_{\alpha}^{i},a_{\alpha}^{i})$
is the generalized $Q$ function.

#### Bias Correction and Control Variates

The method above works well if $\hat{R}(\tau)$ forms an accurate approximation
of $R(\tau)$. However, when the approximation is poor, it introduces a
significant bias into the gradient estimation and hence deteriorates the
performance. To address this problem, we propose to add the residual term to
correct this bias. This yields a gradient estimation of form

$\displaystyle\nabla
J(\theta)=\mathbb{E}_{\pi_{\theta}}\left[r_{0}(\tau)\sum_{t=0}^{T}\nabla\log\pi(a_{t}|s_{t})\right]~{}+~{}\nabla\hat{J}(\theta).$
(7)

where $r_{0}(\tau)=R(\tau)-\hat{R}(\tau)$ denotes the residual error, and
$\nabla\hat{J}(\theta)$ is the gradient of composite reward
$\hat{J}(\theta)=\mathbb{E}_{\pi_{\theta}}[\hat{R}(\tau)]$ in (6) and (7).
This allows us to give an unbiased gradient estimation of the true expected
reward $J(\theta)$, while being able to leverage the structure of the
composite reward.

Theoretically, we can show that the residual corrected gradient estimation (7)
can be viewed as a control variate, a widely used approach for reducing
variance in policy optimization. In particular, we can show that (7) is
equivalent to

$\nabla
J(\theta)=\mathbb{E}_{\pi_{\theta}}\left[\sum_{t=0}^{T}\left(R(\tau)-\hat{r}_{\neg
t}(\tau)\right)\nabla_{\theta}\log\pi(a_{t}|s_{t})\right]$ (8)

where $\hat{r}_{\neg
t}(\tau)=\sum_{\alpha\notin\Gamma_{t}^{*}}\hat{r}(s_{\alpha},a_{\alpha})$,
which collects all the terms that do not appear in (5). We can show that
$\mathbb{E}_{\pi_{\theta}}[\hat{r}_{\neg
t}(\tau)\nabla_{\theta}\log\pi(a_{t}|s_{t})]=0$ for all $t$, and hence
subtracting $r_{\neg t}$ in (8) does not change the expectation of the
formula. Note that $r_{\neg t}$ is a generalization of the standard baseline,
which is typically taken to be a constant, or depend only on $s_{t}$.

With this unbiased estimator, the policy gradient would be robust so that we
no longer need to worry about whether the regression of $\hat{r}$ is
sufficient or the dataset $D$ is not well chosen. In practice, we find that
both $\nabla_{\theta}\hat{J}(\theta)$ and $\nabla_{\theta}J(\theta)$ works
reasonably well if the regression loss is sufficiently optimized.

#### Discussion.

Note that if $\mathcal{I}$ is the set of individual time steps and
$R(\tau)=\sum_{t}\hat{r}(s_{t},a_{t})$, Eq. (8) will be reduced to the classic
policy gradient as shown in Eq. (2). It is probably more interesting to
consider $\mathcal{I}$ to go beyond the set of time steps to include more
information. One critical observation from the above derivation is that the
interval $\alpha$ can be as large as possible if $\max(\alpha)$ is upper
bounded by time $t$, if we want the supervision of
$\hat{r}(s_{\alpha},a_{\alpha})$ for
$\nabla_{\theta}\log\pi_{\theta}(s_{t},a_{t})$. Therefore, to include as much
information as possible, it is natural to see that we can define $\mathcal{I}$
as the set of all consecutive sub-sequences starting from time $0$, so that
each $\alpha\in\\{\\{0,1,2,\ldots,t\\}:\forall t=0,\ldots,T\\}$ contains all
time steps from the beginning of the trajectory to the current time step $t$.
This structure of $\mathcal{I}$ is particularly interesting and highly
resembles the forward-looking structure studied in sequence modeling in NLP,
such as language modeling.

###  2.4 Learning Temporal Credit Assignment using Language Models

Motivated by the forward-looking structure, we propose to adopt neural-network
language models for learning the reward function $\hat{r}$ for credit
assignment. In nature language processing, a language model assigns
probability to a given sequence in a language [5]. A more tangible and related
model is to assign probability of an upcoming word given a sequence of prior
words. More formally, a language model predicts the probability of word
$w_{t}$ by parameterizing the conditional distribution
$p(w_{t}|w_{1},w_{2},\ldots,w_{t-1})$. Such models would be very useful in
many applications, especially those generating sequences as output. Notable
examples including the n-gram models and the recurrent neural networks, which
attempt to capture medium- to long-range dependencies in the sentence. Very
recently, a model entirely based on attention mechanisms was proposed for
language modeling and has achieved state-of-the-art performance on neural
machine translation and other NLP tasks [40, 7]. The model, called
Transformer, has an encoder-decoder structure and is composed of stacked self-
attention and fully connected layers, without using any recurrence or
convolution. To attend multiple parts of the input sequence simultaneously,
instead of using a single large attention “head", the Transformer uses
multiple small attention heads to project the input sequence into multiple
subspaces and combines the attention outputs by concatenation. Note that in
Transformer, the self-attention for constructing an latent vector at a word is
based on all other words prior to the current one, which highly resembles the
desideratum for the reward function, as discussed above. Furthermore, it is
also possible to use recurrent neural networks or convolutional neural
networks to model the reward function.

In this work, we consider the Transformer network model, because of its
superior performance and interpretability. Our model for learning the reward
function $\hat{r}$ consists of two parts. The first part is an encoder module
of the Transformer, which is composed of a multi-head attention layer followed
by a position-wise fully-connected feed-forward layer. The second part is a
self-attention layer [22], which generates a set of summation weight vectors
for the encoder outputs. The set of summation weight vectors are used to
multiply with the Transformer outputs, resulting in a hidden representation
which is then processed by a regression layer to give the predicted rewards
$\hat{r}$.

Specifically, suppose we have a trajectory that has $n$ state-action pairs,
represented as $\tau=\\{(s_{t},a_{t})^{T-1}_{t=0}\\}$. Here $s_{t}$ is a
$d_{s}$-dimensional observation vector and $a_{t}$ is a $d_{a}$-dimensional
action vector. $\tau$ is thus represented as a $T$ by $(d_{s}+d_{a})$ matrix.
Note that we omit the special terminal state $s_{T}$ here for notational
simplicity, if not causing further confusion. Each state-action pair
$(s_{t},a_{t})$ in the trajectory is then processed by a feed forward layer,
whose parameters are shared across all time steps, to give a fixed-length
vector representation $\mathbf{v}_{t}$ for this state-action pair. To gain the
dependency between the current time step $t$ and other time steps before $t$,
we use a encoder layer to process the state-action pairs:
$\mathbf{h}_{t}=\text{Transformer}(\mathbf{v}_{0},\mathbf{v}_{1},\ldots,\mathbf{v}_{t})$.
Let the dimension of each start-action pair representation vector
$\mathbf{v}_{t}$ be $d$. Since the Transformer network does not change the
dimension of input vectors, we represent all the $T$ representation vectors
$\mathbf{h}_{t}$ as a $T$ by $d$ matrix
$H=(\mathbf{h}_{0},\mathbf{h}_{1},\ldots,\mathbf{h}_{T-1})$. To summarize the
information in $H$, we apply a self-attention mechanism:

$\mathbf{z}=\text{sigmoid}(\mathbf{w}_{s2}\tanh(W_{s1}H^{\top}))^{\top}.$

Here $W_{s1}$ is a weight matrix with dimension $d_{z}$ by $d$ and
$\mathbf{w}_{s2}$ is a vector of parameters with size $d_{z}$, where $d_{z}$
is a hyper-parameter. Vector $\mathbf{z}$ has size $T$, and each entry $z_{t}$
ranges from 0 to 1, quantifying the importance of the state-action pair
$(s_{t},a_{t})$ in predicting the reward $r(s_{\alpha_{t}},a_{\alpha_{t}})$
for interval $\alpha_{t}=\\{0,1,\ldots,t\\}$. We combine the hidden
representations in $H$ using $\mathbf{z}$ to obtain a summarized
representation $\mathbf{h}^{*}_{t}=z_{t}\mathbf{h}_{t}$. To predict the reward
$\hat{r}(s_{\alpha_{t}},a_{\alpha_{t}})$, we add one regression layer
parameterized by $\mathbf{w}_{r}$ and $b_{r}$ and output the predicted reward
as

$\hat{r}(s_{\alpha_{t}},a_{\alpha_{t}})=\mathbf{w}_{r}^{\top}\mathbf{h^{*}_{t}}+b_{r}.$

Some other networks we consider include the feed-forward neural network and
long short-term memory network (LSTM). Please see Figure A2 to see their
differences.

##  3 Experiments

In this section, we conduct experiments to provide evidences for the following
questions? (1) Is the learned reward function useful for improving policy
optimization in episodic RL? (2) What are the appropriate choices of interval
set $\mathcal{I}$, neural network model, and dataset $D$? Note that we also
provide a case study on visualization and intepretability of the learned
reward function in Appendix.

###  3.1 Experimental Settings

The experiments are conducted on a set of high-dimensional locomotion tasks in
continuous domain using OpenAI gym [6] and MuJoCo simulation toolkits. We use
the PPO algorithm introduced in [32] for all the experiments. The policy is
represented by a uni-mode Gaussian distribution with diagonal covariance. The
mean is parameterized by a two-layer neural network with 64 hidden units and
tanh non-linearity. The log standard deviation is parameterized by a global
vector. The same architecture is applied for value function approximation. In
practice, we implemented the policy gradient estimation shown in Eq. (7),
where we decompose into two parts: one with the predicted reward, the other
with only the residual. In this way, we can use advanced variance reduction
approaches to improve the estimation of each components. In our experiments,
we use the generalized advantage estimation (GAE) [31] as the control variate
method for variance reduction. More hyper-parameters of PPO and GAE are
tabulated in Table A1 in Appendix. We also compare against baseline algorithms
trained with episodic reward, respectively. The episodic reward of each
rollout trajectory is defined as the accumulated original reward at all time-
steps. The experiments are run for 5M timesteps with 5 random seeds. For
baseline algorithms, we also consider an LSTM policy with 128 hidden units
besides the aforementioned MLP policy. All experiments are conducted on NVIDIA
1080 GPUs.

| Hopper | Walker2d | Humanoid | Humanoid-Standup | Swimmer
---|---|---|---|---|---
PPO (episodic) | 437 | 266 | 516 | 44673 | 6
CEM | 97 | 205 | 426 | $\approx 9.6\times 10^{4}$ | 17
Ours | 1462 | 3217 | 2209 | 82579 | 135
Table 1: Performance of PPO and our approach on locomotion tasks with episodic
rewards. Scores taken at 5M iterations with the environment. Cross entropy
method performance taken from [8].

For the reward function in our experiments, we investigate three neural
network architectures, Feed-Forward (FF) Network, LSTM [13], and Transformer,
to identify the best network structure that captures temporal dependency
(Figure A2). For observation and action $s_{t},a_{t}$ at time $t$, an FF
network, whose parameter $\phi$ is shared across all time steps, is used to
give the predicted reward $\hat{r}_{t}$ for individual time steps. LSTM’s
reward prediction is also conditioned on all previous timesteps, so is a
function of trajectory segment $s_{0:t},a_{0:t}$. To perform reward regression
based on the LSTM/FF outputs, we aggregate the the representations using a
mean-pooling layer. The hyper-parameters of LSTM and FF are also listed in A1.

For the buffer updating schemes, we proposed the following three approaches to
renew the trajectory buffer for reward predicting.

  * •

Online (O): The buffer is implemented as a FIFO queue with length $K$ (hyper-
parameter). In each iteration, the new roll-out trajectories will be inserted
into the queue and the old trajectories will be removed if the queue is too
long.

  * •

Historical+Online (HO): Two queues are maintained: One queue is the same as
the one in Online, storing the most recent rollouts by the current policy. The
other queue stores the trajectories with the highest episodic returns from
previous rollouts.

  * •

Stratified-Sampling (S): To balance the training with episodic return
regression, the buffer stores a larger number of trajectories in the history.
In each iteration, $K$ trajectories are sampled to ensure the episodic return
is uniformly in the five bins. The queue size $L$ and sample queue size $K$
are hyper-parameters.

###  3.2 Credit Assignment Enables Policy Optimization with Episodic Reward

Our experiments with the baseline algorithms on the MuJoCo control suite
demonstrated that learning with episodic return is extremely hard both with
MLP policy and LSTM policy. In all the environments we tested (Figure 1), it
is not surprising that the baseline method PPO (episodic) with episodic reward
performs much worse than previously reported in other papers with dense reward
[32]. In most of the runs, the policy cannot make any improvement after the
initial time steps.

Figure 1: Learning curves for PPO baselines and our proposed method on tasks
with episodic rewards. Mean and standard deviation over 5 random seeds is
plotted. The x- and y- axis represent the number of training samples (in
million) and average return, respectively.

On the other hand, our proposed credit-assignment algorithm with learned
return decomposition consistently achieves better performance than the
episodic return baselines across all environments (Figure 1). Here we use the
Transformer network for decomposing rewards and MLP for representing policy,
and the HO strategy to collect trajectories for optimization. In many
experiments, the policies learned by our method are able to achieve the quite
reasonably good performance when the original dense reward is used for
training. In environments such as the Humanoid, the Hopper and the Swimmer,
with the appropriate hyper-parameters, our methods can obtain comparable
performance to the policies trained with the original dense rewards,
outperforming the episodic return baseline by a large margin. In addition, it
also outperforms the cross-entropy method (CEM) which is suitable for the
episodic RL setting (Table 1). These results indicate that the our reward
decomposition framework can successfully enable stable learning and greatly
improve the sample efficiency in episodic settings.

###  3.3 Ablation Analysis

Here we perform comparisons to check the choices of 1) three network
structures for reward function; 2) strategies for data collection. We also
check the utility of bias correction when the reward regression is not well
fitted. We use the environment Walker2d to perform these analyses and
demonstrate the results in Figure 2.

As we can see, all the three networks structures outperform the episodic
baselines, and the Transformer network provides the best performance. We
conjecture that the reason why Transformer performs better than LSTM is that
it is easier to train, as also implied by existing NLP literature [40]. Then
we check whether the data collection strategy has any impact on the
performance. All three strategies seem to work reasonably well, while HO
appears to be the best which indicating that the historical high-quality
trajectories helps reward learning. Finally, we compare the performance of
$\nabla_{\theta}J(\theta)$ and $\nabla_{\theta}\hat{J}(\theta)$ under
different learning rates ($10^{-2}$ and $10^{-3}$) for the regression loss. We
use different learning rates essentially to adjust the quality of the reward
fitting. It is easy to see that with the bias correction, the learning of
$\nabla_{\theta}J(\theta)$ is more robust, indicating that bias correction is
needed. Extra results on ablation analysis can be found in the appendix
(Figures A3, A4).

Figure 2: Ablation analysis on choices of (a) network structures for reward
function; (b) strategies for data collection; (c)-(d) methods with or without
bias correction under different learning rates (lr). Environment Walker2d was
used to perform the analyses. Mean and standard deviation over 5 random seeds
is plotted.

##  4 Related Work

The optimal reward design problem [35, 36] concerns about finding the proxy
reward function can obtains high expected return according to the true reward
function. Inverse reward design [10] studies the opposite problem of inferring
true reward from designed reward. Intrinsic motivation has been shown to
improve sample efficiency of RL algorithms, for example by using information
gain [15], pseudo count [4] or prediction error [37, 28] as an intrinsic bonus
reward to aid exploration. [42] studies intrinsic motivation under the optimal
reward design framework where the optimal intrinsic reward is learned through
gradient descent. Reward shaping [26] explores the space of reward function
modifications (specifically potential-based rewards) which do not change the
corresponding optimal policy. Hindsight Experience Replay [2] adds additional
goals and corresponding rewards to a Q-learning algorithm. Meta-learning can
also be utilized to learn different objective functions for different goals
[14]. Our approach is more general as we don’t assume specific goals of the
agent. When only expert demonstrations are available as in imitation learning,
inverse reinforcement learning can be used to recover the reward function from
expert trajectories [1]. As an alternative approach for solving sparse reward
problems, the auxiliary-task approaches [17] focus on improving the
representation by adding extra self-supervised losses. However, our proposed
method learns to decompose the episodic return as reward for policy
optimization directly.

Recently, several concurrent works have explored the same direction of
decomposing the episodic return. Sparse Attentive Backtracking[19] applies the
attentive mechanism for network straining, while ours focuses on learning the
decomposition of the reward signal. Temporal Value Transport (TVT) [16] relies
on a memory reconstitutive module to retrieve past data for the visual RL
problems, however, ours aims at decomposing the episodic return for general
problems. As one most related work, RUDDER [3] could be viewed as a special
case of our work in which consecutive time-steps were used. On the other hand,
our work provides a general framework with interval rewards which justifies
the applicability of language models and the correctness of generalized policy
gradient beyond the simple Markovian assumption.

##  5 Conclusion

We presented a new algorithm for learning temporal credit assignment, which
uses deep neural network (Transformer) to decompose the episodic reward back
to each time step in the trajectory. The assigned dense reward signals
obtained from the decomposition are then used to guide training algorithms of
policies. We demonstrated that our credit assignment algorithm substantially
improved the learning and sample efficiency on a set of MuJoCo continuous
locomotive control tasks. The reward function learned by our algorithm can
also be interpreted by an attention mechanism, potentially providing insights
on identifying key state-action pairs that contribute to successful
reinforcement learning.

## References

  * [1] Pieter Abbeel and Andrew Y Ng.  Apprenticeship learning via inverse reinforcement learning.  In Proceedings of the twenty-first international conference on Machine learning, page 1. ACM, 2004.
  * [2] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAI Pieter Abbeel, and Wojciech Zaremba.  Hindsight experience replay.  In Advances in Neural Information Processing Systems, pages 5048–5058, 2017.
  * [3] Jose A Arjona-Medina, Michael Gillhofer, Michael Widrich, Thomas Unterthiner, and Sepp Hochreiter.  Rudder: Return decomposition for delayed rewards.  arXiv preprint arXiv:1806.07857, 2018.
  * [4] Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos.  Unifying count-based exploration and intrinsic motivation.  In Advances in Neural Information Processing Systems, pages 1471–1479, 2016.
  * [5] Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Jauvin.  A neural probabilistic language model.  Journal of machine learning research, 3(Feb):1137–1155, 2003.
  * [6] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym.  arXiv preprint arXiv:1606.01540, 2016.
  * [7] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  arXiv preprint arXiv:1810.04805, 2018.
  * [8] Tanmay Gangwani, Qiang Liu, and Jian Peng.  Learning self-imitating diverse policies.  Proceedings of the International Conference on Learning Representations (ICLR), 2019.
  * [9] Yijie Guo, Junhyuk Oh, Satinder Singh, and Honglak Lee.  Generative adversarial self-imitation learning.  arXiv preprint arXiv:1812.00950, 2018.
  * [10] Dylan Hadfield-Menell, Smitha Milli, Pieter Abbeel, Stuart J Russell, and Anca Dragan.  Inverse reward design.  In Advances in Neural Information Processing Systems, pages 6765–6774, 2017.
  * [11] Nikolaus Hansen and Andreas Ostermeier.  Completely derandomized self-adaptation in evolution strategies.  Evolutionary computation, 9(2):159–195, 2001.
  * [12] Nicolas Heess, Srinivasan Sriram, Jay Lemmon, Josh Merel, Greg Wayne, Yuval Tassa, Tom Erez, Ziyu Wang, Ali Eslami, Martin Riedmiller, et al.  Emergence of locomotion behaviours in rich environments.  arXiv preprint arXiv:1707.02286, 2017.
  * [13] Sepp Hochreiter and Jürgen Schmidhuber.  Long short-term memory.  Neural computation, 9(8):1735–1780, 1997.
  * [14] Rein Houthooft, Richard Y Chen, Phillip Isola, Bradly C Stadie, Filip Wolski, Jonathan Ho, and Pieter Abbeel.  Evolved policy gradients.  arXiv preprint arXiv:1802.04821, 2018.
  * [15] Rein Houthooft, Xi Chen, Yan Duan, John Schulman, Filip De Turck, and Pieter Abbeel.  Vime: Variational information maximizing exploration.  In Advances in Neural Information Processing Systems, pages 1109–1117, 2016.
  * [16] Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne.  Optimizing agent behavior over long time scales by transporting value.  arXiv preprint arXiv:1810.06721, 2018.
  * [17] Max Jaderberg, Volodymyr Mnih, Wojciech Czarnecki, Tom Schaul, Joel Z. Leibo, David Silver, and Koray Kavukcuoglu.  Reinforcement learning with unsupervised auxiliary tasks.  CoRR, abs/1611.05397, 2017.
  * [18] Daniel Jurafsky and James H. Martin.  Speech and Language Processing (2Nd Edition).  Prentice-Hall, Inc., Upper Saddle River, NJ, USA, 2009.
  * [19] Nan Rosemary Ke, Anirudh Goyal, Olexa Bilaniuk, Jonathan Binas, Michael C. Mozer, Chris Pal, and Yoshua Bengio.  Sparse attentive backtracking: Temporal credit assignment through reminding.  In Proceedings of the 32Nd International Conference on Neural Information Processing Systems, NIPS’18, pages 7651–7662, USA, 2018. Curran Associates Inc.
  * [20] Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel.  End-to-end training of deep visuomotor policies.  The Journal of Machine Learning Research, 17(1):1334–1373, 2016\.
  * [21] Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra.  Continuous control with deep reinforcement learning.  ICLR, 2016.
  * [22] Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio.  A structured self-attentive sentence embedding.  arXiv preprint arXiv:1703.03130, 2017.
  * [23] Hongzi Mao, Mohammad Alizadeh, Ishai Menache, and Srikanth Kandula.  Resource management with deep reinforcement learning.  In Proceedings of the 15th ACM Workshop on Hot Topics in Networks, pages 50–56. ACM, 2016.
  * [24] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu.  Asynchronous methods for deep reinforcement learning.  In International conference on machine learning, pages 1928–1937, 2016.
  * [25] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  Nature, 518(7540):529, 2015.
  * [26] Andrew Y Ng, Daishi Harada, and Stuart Russell.  Policy invariance under reward transformations: Theory and application to reward shaping.  In ICML, volume 99, pages 278–287, 1999.
  * [27] Marcus Olivecrona, Thomas Blaschke, Ola Engkvist, and Hongming Chen.  Molecular de-novo design through deep reinforcement learning.  Journal of cheminformatics, 9(1):48, 2017.
  * [28] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  In International Conference on Machine Learning (ICML), volume 2017, 2017.
  * [29] Reuven Rubinstein.  The cross-entropy method for combinatorial and continuous optimization.  Methodology and computing in applied probability, 1(2):127–190, 1999.
  * [30] Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever.  Evolution strategies as a scalable alternative to reinforcement learning.  arXiv preprint arXiv:1703.03864, 2017.
  * [31] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel.  High-dimensional continuous control using generalized advantage estimation.  Proceedings of the International Conference on Learning Representations (ICLR), 2016.
  * [32] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  arXiv preprint arXiv:1707.06347, 2017.
  * [33] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al.  Mastering the game of go with deep neural networks and tree search.  nature, 529(7587):484, 2016.
  * [34] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al.  Mastering the game of go without human knowledge.  Nature, 550(7676):354, 2017.
  * [35] Satinder Singh, Richard L Lewis, and Andrew G Barto.  Where do rewards come from.  In Proceedings of the annual conference of the cognitive science society, pages 2601–2606, 2009.
  * [36] Jonathan Sorg, Richard L Lewis, and Satinder P Singh.  Reward design via online gradient ascent.  In Advances in Neural Information Processing Systems, pages 2190–2198, 2010.
  * [37] Bradly C Stadie, Sergey Levine, and Pieter Abbeel.  Incentivizing exploration in reinforcement learning with deep predictive models.  arXiv preprint arXiv:1507.00814, 2015.
  * [38] Richard S Sutton and Andrew G Barto.  Reinforcement learning: An introduction.  MIT press, 2018.
  * [39] Richard Stuart Sutton.  Temporal Credit Assignment in Reinforcement Learning.  PhD thesis, University of Massachusetts Amherst, 1984.  AAI8410337.
  * [40] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In Advances in Neural Information Processing Systems, pages 5998–6008, 2017.
  * [41] Ronald J Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  Machine learning, 8(3-4):229–256, 1992.
  * [42] Zeyu Zheng, Junhyuk Oh, and Satinder Singh.  On learning intrinsic rewards for policy gradient methods.  In Advances in Neural Information Processing Systems, pages 4649–4659, 2018.

##  Appendix A Method Overview and Network structures

Figure A1:  Overview of our approach. Rollout trajectories are generated from
interacting with the environment. A reward predictor is trained on the
collected trajectories and episodic returns with the regression loss. Then the
predicted rewards are used for policy optimization. Figure A2:  Network
structures for the reward predictor: (a) Feed-forward network, (b) LSTM
network, (c) Transformer network

##  Appendix B Proof of Theorem 1

I) Recall the standard likelihood ratio gradient formula:

$\displaystyle\nabla_{\theta}\hat{J}(\theta)$
$\displaystyle=\nabla_{\theta}\mathbb{E}_{\pi_{\theta}}[\hat{R}(\tau)]$
$\displaystyle=\mathbb{E}_{\pi_{\theta}}\left[\hat{R}(\tau)\sum_{t=0}^{T}\nabla_{\theta}\log\pi(a_{t}|s_{t})\right]$
$\displaystyle=\mathbb{E}_{\pi_{\theta}}\left[\left(\sum_{\alpha\in\mathcal{I}}\hat{r}(s_{\alpha},a_{\alpha})\right)\left(\sum_{t=0}^{T}\nabla_{\theta}\log\pi(a_{t}|s_{t})\right)\right]$
$\displaystyle=\sum_{\alpha\in\mathcal{I}}\mathbb{E}_{\pi_{\theta}}\left[\hat{r}(s_{\alpha},a_{\alpha})\sum_{t=0}^{T}\nabla_{\theta}\log\pi(a_{t}|s_{t})\right].$
(9)

On the other hand, note that for any $t\notin\Gamma_{\alpha}$, we have

$\displaystyle\mathbb{E}_{\pi_{\theta}}[\hat{r}(s_{\alpha},s_{\alpha})\nabla_{\theta}\log\pi(a_{t}|s_{t})]=0.$

Therefore, all the pairs $(\alpha,t)$ with $t\notin\Gamma_{\alpha}$ is removed
in (9). This hence yields (5).

II) Eq. (6) is a simple rearrangement of Eq. (5).

$\displaystyle\nabla_{\theta}\hat{J}(\theta)$
$\displaystyle=\sum_{\alpha\in\mathcal{I}}\mathbb{E}_{\pi_{\theta}}\left[\hat{r}(s_{\alpha},a_{\alpha})\sum_{t\in\Gamma_{\alpha}}\nabla_{\theta}\log\pi(a_{t}|s_{t})\right]$
$\displaystyle=\mathbb{E}_{\pi_{\theta}}\left[\sum_{t=0}^{T}\left(\sum_{\alpha\colon
t\ni\Gamma_{\alpha}}\hat{r}(s_{\alpha},a_{\alpha})\right)\nabla_{\theta}\log\pi(a_{t}|s_{t})\right],$

where $(\sum_{\alpha\colon
t\ni\Gamma_{\alpha}}\hat{r}(s_{\alpha},a_{\alpha})=Q_{t}(\tau)$, matching our
definition. This completes the proof.

##  Appendix C Hyper-parameters

Hyper-parameter | Searching Values
---|---
Policy Network | MLP with shape (64, 64)
PPO Batch Size | 2048
PPO Mini-Batch Size | 64
PPO number of epoch per iteration | 5
PPO learning rate | 0.0001
PPO clip range $\epsilon$ | 0.2
GAE $\gamma$ | 0.99
GAE $\lambda$ | 0.95
Buffer size | 50
Reward predictor learning rate | $0.001$
Transformer number of heads | 4
Transformer layer size | 64
Transformer hidden layer size | 128
Transformer query/key size | 32
LSTM hidden size | 96
FF channels | $\left[128,128,128,256\right]$
Table A1: The hyper-parameters we used in the episodic MuJoCo environment.

##  Appendix D More Results on Ablation Analysis

Figure A3: Comparison between different buffer updating methods. The x-axis
denotes the number of training samples and y-axis denotes the average episodic
return. The red curve represents the training curve using episodic return. The
yellow, blue and green curves represent the algorithm with online buffer
scheme, historical-online scheme and stratified-sampling scheme, respectively.
Figure A4: Comparison between network structures: The x- and y- axis represent
the number of training samples and average return, respectively. The red curve
represents the training curve using episodic return. The blue, yellow and
green lines represent our method with a Transformer, LSTM and FF network
structure, respectively.

##  Appendix E Interpretability of the Learned Reward Function

We visualize the learned agent and reward function to demonstrate that the
reward predictor could attain knowledge from interacting with environment. Key
state action pairs that contribute to successful reinforcement learning can
also identified from the visualization . Using the Hopper environment as an
example, we visualized in Figure A5 the learned temporal attentions for 1000
time steps, extracted from the last layer of the Transformer network. The
first stage A corresponds to the agent starting a large jump, followed by the
landing (stage B) and laying on the ground (stage C). The agent then makes a
smaller hop in stage D and lands in stage E. We observed that the temporal
attentions exhibit periodic behavior that coincides with the agent’s periodic
motion. In the Hopper environment, the goal is to make the 2D one-legged robot
move forward as fast as possible. A successfully trained agent essentially
learns to hop forward periodically. Hence it is reasonable that the temporal
attention and predicted reward show periodicity as the hopper jumps. Similar
periodic behaviors of moving and adjusting balance were also be observed in
other locomotion environments such as Humanoid.

Figure A5:  (Above) Learned temporal attention in the Transformer structure
shows periodicity over time in the Hopper environment. A single period
consists of five stages as marked by A to E. The model gives higher attention
to jumping (A, D) than landing (B, E). (Below) visualizes the learned state
dependency, i.e. the key reference states (e.g. landing at t=130) predicted by
the multi-head attention layer in Transformer for states (e.g. jumping at
t=495, 500, 505). We observed that the learned reward function gives higher
reward and attention to the agent’s jumping, as the rewards of both the
jumping in the large hop (stage A) and the small hop (stage D) are relatively
higher than that of landing phases (stages B, C and E).