11institutetext:  I.E. Tamm Department of Theoretical Physics, P.N. Lebedev
Physical Institute, 119991 Moscow, Russia 22institutetext: Institute for
Quantum Materials and Technologies, Karlsruhe Institute of Technology (KIT),
76021 Karlsruhe, Germany 33institutetext: National Research University Higher
School of Economics, 101000 Moscow, Russia

# Long-range Josephson effect controlled by temperature gradient and circuit
topology

Mikhail S. Kalenkov 11 Andrei D. Zaikin 2233<EMAIL_ADDRESS>
###### Abstract

We demonstrate that the supercurrent can be strongly enhanced in cross-like
superconducting hybrid nanostructures ($X$-junctions) exposed to a temperature
gradient. At temperatures $T$ exceeding the Thouless energy of our
$X$-junction the Josephson current decays algebraically with increasing $T$
and can be further enhanced by a proper choice of the circuit topology. At
large values of the temperature gradient the non-equilibrium contribution to
the supercurrent may become as large as the equilibrium one at low $T$. We
also predict a variety of transitions between 0- and $\pi$-junction states
controlled by the temperature gradient as well as by the system geometry. Our
predictions can be directly verified in modern experiments.

##  1 Introduction

Equilibrium Josephson current between two superconductors depends periodically
on the phase difference $\chi$ between them Jos . The magnitude of this
current also depends on temperature $T$. This dependence varies in different
types of superconducting weak links and is perhaps most strongly pronounced in
superconducting junctions containing a sufficiently thick layer of a normal
metal ZZh ; GreKa ; BWBSZ1999 ; GKI . In these so-called SNS junctions the
supercurrent reaches its maximum value $I_{C}(0)\simeq
10.82E_{\mathrm{Th}}/(eR_{n})$ GreKa  at $T\to 0$, whereas at higher
temperatures $T>E_{\mathrm{Th}}$ this current reduces exponentially as
$\propto e^{-\sqrt{2\pi T/E_{\mathrm{Th}}}}$, where $E_{\mathrm{Th}}$ and
$R_{n}$ are respectively an effective Thouless energy and a normal state
resistance of an SNS device.

The supercurrent flowing across an SNS junction can be significantly affected
by driving the electron distribution function out of equilibrium. Such non-
equilibrium conditions can be achieved, e.g., by applying an external ac
signal Aslamazov82 ; Zaikin83  or a dc voltage $V$ in the cross-like geometry
considered in Volkov ; WSZ ; Yip ; Teun . In the first case one can observe a
strong supercurrent stimulation at $T>E_{\mathrm{Th}}$, while in the second
one by tuning $V$ one can realize the transition to a $\pi$-junction state.

Yet another way to drive the electron distribution function inside an SNS
junction out of equilibrium is to expose it to a thermal gradient. Recently it
was demonstrated KDZ20  that by doing so one can effectively support long-
range phase coherence of quasiparticles inside the N-layer at temperatures
above the Thouless energy where the equilibrium supercurrent already becomes
vanishingly small. Specifically, for the so-called $X$-junction geometry
illustrated in Fig. 1 in the high temperature limit $T_{1,2}\gg
E_{\mathrm{Th}}$ – up to some geometry factors – one finds $I_{C}\sim
I_{C}(0)E_{\mathrm{Th}}|1/T_{1}-1/T_{2}|$, where $T_{1}$ and $T_{2}$ are
different temperatures at which two normal terminals N1 and N2 are maintained,
see Fig. 1. Hence, this current turns out to be a lot bigger than the
equilibrium one at any of the two temperatures $T_{1}$ or $T_{2}$.

Furthermore, under such non-equilibrium conditions the junction is described
by a non-sinusoidal current-phase relation and may exhibit a pronounced
$\pi$-junction-like behavior. It was argued KDZ20  that all these non-trivial
features are caused by the presence of non-equilibrium low energy
quasiparticles suffering little dephasing while propagating across the
N-layer. We also note that a somewhat similar situation was encountered for
the Aharonov-Bohm effect in superconducting heterostructures containing a
normal metallic loop, see, e.g., Refs. BWBSZ1999 ; Gre ; GWZ .

In this paper we will extend the work KDZ20  in several important aspects. In
particular, here we will lift the symmetry restrictions adopted in KDZ20  and
evaluate the Josephson current across a general asymmetric $X$-junction
exposed to an arbitrary temperature gradient. We will demonstrate that, on one
hand, the presence of electron-hole asymmetry weakly affects the Josephson
current and, on the other hand, that the effect of supercurrent stimulation
can be substantially enhanced by a proper choice of the circuit topology.

Figure 1: $X$-junction structure under consideration.

##  2 The model and basic formalism

Following Ref. KDZ20  we are going to consider an $X$-junction which consists
of two superconducting and two normal terminals interconnected by four normal
metallic wires of lengths $L_{S_{1,2}}$, $L_{N_{1,2}}$ and cross sections
${\mathcal{A}}_{S_{1,2}}$, ${\mathcal{A}}_{N_{1,2}}$ respectively, as shown in
Fig. 1. The two normal terminals are disconnected from any external circuit
and are maintained at different temperatures $T_{1}$ and $T_{2}$. The
supercurrent $I_{S}(\chi)$ can flow between superconducting terminals S1 and
S2 with the phase difference $\chi=\chi_{1}-\chi_{2}$.

In what follows we will employ the standard quasiclassical formalism based on
the Usadel equations combined with the Keldysh technique BWBSZ1999

$iD\nabla\left(\check{G}\nabla\check{G}\right)=\left[\check{\Omega},\check{G}\right],\quad\check{G}\check{G}=\check{1},$
(1)

where $D$ is the diffusion constant, $\check{\Omega}$ is $4\times 4$ matrix

$\check{\Omega}=\begin{pmatrix}\hat{\Omega}^{R}&0\\\
0&\hat{\Omega}^{A}\end{pmatrix},\quad\hat{\Omega}^{R}=\hat{\Omega}^{A}=\begin{pmatrix}\varepsilon+eV&\Delta\\\
-\Delta^{*}&-\varepsilon+eV\end{pmatrix},$ (2)

where $\varepsilon$, $V$ and $\Delta$ denote respectively quasiparticle
energy, electrostatic potential and superconducting order parameter. Equations
(1) allow to evaluate the quasiclassical Green function $\check{G}$, which is
represented by the $4\times 4$ matrix in the Keldysh$\otimes$Nambu space

$\check{G}=\begin{pmatrix}\hat{G}^{R}&\hat{G}^{K}\\\
0&\hat{G}^{A}\end{pmatrix},\quad\hat{G}^{R,A}=\begin{pmatrix}G^{R,A}&F^{R,A}\\\
\tilde{F}^{R,A}&-G^{R,A}\end{pmatrix},$ (3)

where $\hat{G}^{R,A}$ are retarded and advanced Green functions. It is
convenient to parametrize the Keldysh matrix function $\hat{G}^{K}$ as

$\hat{G}^{K}=\hat{G}^{R}\hat{h}-\hat{h}\hat{G}^{A},\quad\hat{h}=h^{L}+\hat{\tau}_{3}h^{T},$
(4)

where $h^{L}$ and $h^{T}$ are the parts of the quasiparticle distribution
function. The current density can be calculated by means of standard relation

$\displaystyle\bm{j}=-\dfrac{\sigma}{8e}\int
d\varepsilon\operatorname{\mathrm{Sp}}(\hat{\tau}_{3}\check{G}\nabla\check{G})^{K}.$
(5)

Inside the normal wires the functions $h^{L}$ and $h^{T}$ obey the equations

$\displaystyle iD\nabla\left[D^{T}\nabla h^{T}+{\mathcal{Y}}\nabla
h^{L}+\bm{j}_{\varepsilon}h^{L}\right]=0,$ (6) $\displaystyle
iD\nabla\left[D^{L}\nabla h^{L}-{\mathcal{Y}}\nabla
h^{T}+\bm{j}_{\varepsilon}h^{T}\right]=0,$ (7)

which follow directly from Eqs. (1). Here

$\displaystyle
D^{T}=\dfrac{1}{4}\operatorname{\mathrm{Sp}}(1-\hat{\tau}_{3}\hat{G}^{R}\hat{\tau}_{3}\hat{G}^{A})=\dfrac{1}{4}\left[2-2G^{R}G^{A}+F^{R}\tilde{F}^{A}+\tilde{F}^{R}F^{A}\right],$
(8) $\displaystyle
D^{L}=\dfrac{1}{4}\operatorname{\mathrm{Sp}}(1-\hat{G}^{R}\hat{G}^{A})=\dfrac{1}{4}\left[2-2G^{R}G^{A}-F^{R}\tilde{F}^{A}-\tilde{F}^{R}F^{A}\right]$
(9)

denote dimensionless diffusion coefficients,

$\mathcal{Y}=\dfrac{1}{4}\left[\tilde{F}^{R}F^{A}-F^{R}\tilde{F}^{A}\right]$
(10)

is the kinetic coefficient which accounts for the presence of the particle-
hole asymmetry in our system and

$\bm{j}_{\varepsilon}=\dfrac{1}{2}\operatorname{\mathrm{Re}}\left(F^{R}\nabla\tilde{F}^{R}-\tilde{F}^{R}\nabla
F^{R}\right)$ (11)

defines the spectral supercurrent.

##  3 General analysis

For quasi-one-dimensional geometry of the normal wires adopted here the
kinetic equations (6), (7) can be solved exactly. Let us first rewrite these
equations in the matrix form

$\displaystyle\hat{D}\begin{pmatrix}h^{T}\\\
h^{L}\end{pmatrix}^{\prime}+j_{\varepsilon}\hat{\tau}_{1}\begin{pmatrix}h^{T}\\\
h^{L}\end{pmatrix}=-\dfrac{e}{\sigma\mathcal{A}}\begin{pmatrix}I^{T}\\\
I^{L}\end{pmatrix},\quad\hat{D}=\begin{pmatrix}D^{T}&\mathcal{Y}\\\
-\mathcal{Y}&D^{L}\end{pmatrix}$ (12)

where $\mathcal{A}$ is the cross section of the corresponding wire segment.
The spectral currents $I^{T}$ and $I^{L}$ represent the integration constants
independent of the coordinate along the wire. Below we will use the convention
according to which the current is positive if it flows from the corresponding
terminal to the crossing point $c$.

As the spectral supercurrent $j_{\varepsilon}$ vanishes identically in the
wires attached to the normal terminals, Eq. (12) can be easily integrated, and
we get

$\hat{G}_{N_{i}}\begin{pmatrix}h^{T}_{c}-h^{T}_{N_{i}}\\\
h^{L}_{c}-h^{L}_{N_{i}}\end{pmatrix}=\begin{pmatrix}I^{T}_{N_{i}}\\\
I^{L}_{N_{i}}\end{pmatrix},\quad\hat{G}_{N_{i}}=\begin{pmatrix}G^{T}_{N_{i}}&G^{\mathcal{Y}}_{N_{i}}\\\
-G^{\mathcal{Y}}_{N_{i}}&G^{L}_{N_{i}}\end{pmatrix}=\left[\int_{L_{N_{i}}}\dfrac{\hat{D}^{-1}dx}{\sigma\mathcal{A}_{N_{i}}}\right]^{-1}.$
(13)

The spectral conductances $G^{T}_{N_{i}}$, $G^{L}_{N_{i}}$ and
$G^{\mathcal{Y}}_{N_{i}}$ all exhibit a nontrivial energy dependence in the
vicinity of the Thouless energy, whereas for $|\varepsilon|\gg
E_{\mathrm{Th}}$ the conductance $G^{\mathcal{Y}}_{N_{i}}$ tends to zero and
$G^{T}_{N_{i}}$, $G^{L}_{N_{i}}$ just reduce to their normal state values
$G^{n}_{N_{i}}\equiv 1/R^{n}_{N_{i}}=\sigma\mathcal{A}_{N_{i}}/L_{N_{i}}$.

In the wires attached to the superconducting terminals and at subgap energies
the spectral currents $I^{L}_{S_{i}}$ vanish identically. With this in mind,
we obtain

$\displaystyle
G^{T}_{S_{1}}h_{c}^{T}+\mathcal{G}h_{c}^{L}=-eI^{T}_{S_{1}},\quad
G^{T}_{S_{2}}h_{c}^{T}-\mathcal{G}h_{c}^{L}=-eI^{T}_{S_{2}},$ (14)

where

$\mathcal{G}(\varepsilon)=\mathcal{A}_{S_{1}}\sigma
j_{1\varepsilon}=-\mathcal{A}_{S_{2}}\sigma j_{2\varepsilon}.$ (15)

Here we also made use of the condition $h^{T}=0$ which holds for both SN
interfaces at subgap energies. In general the spectral conductances
$G^{T}_{S_{1,2}}$ depend on the kinetic coefficients $D^{T,L}$ and
$\mathcal{Y}$ in a complicated manner. These conductances demonstrate a
nontrivial energy dependence below the Thouless energy and tend to their
normal state values $G^{n}_{S_{1,2}}=\sigma\mathcal{A}_{S_{1,2}}/L_{S_{1,2}}$
in the high energy limit.

The spectral current conservation conditions at the crossing point $c$ take
the form

$\displaystyle I^{T}_{S_{1}}+I^{T}_{S_{2}}+I^{T}_{N_{1}}+I^{T}_{N_{2}}=0,\quad
I^{L}_{N_{1}}+I^{L}_{N_{2}}=0.$ (16)

Resolving the above equations we can express all the spectral currents
$I^{T,L}$ in terms of the distribution functions in the normal terminals

$h^{T/L}_{N_{1,2}}=\dfrac{1}{2}\left[\tanh\dfrac{\varepsilon+eV_{1,2}}{2T_{1,2}}\mp\tanh\dfrac{\varepsilon-
eV_{1,2}}{2T_{1,2}}\right].$ (17)

As soon as the spectral currents are established, the electric currents
$I_{X_{i}}$ in all four normal wires can be recovered with the aid of a simple
formula

$I_{X_{i}}=\dfrac{1}{2}\int I_{X_{i}}^{T}d\varepsilon,\quad X=S,N,\quad
i=1,2.$ (18)

##  4 Supercurrent in the presence of a temperature gradient

Without any temperature gradient no voltage drop across the normal terminals
N1 and N2 can occur. However, different electrostatic potentials $V_{1}$ and
$V_{2}$ at these terminals are in general induced provided they are kept at
different temperatures $T_{1}$ and $T_{2}$. This is a manifestation of the so-
called thermoelectric effect Ginzburg . In order to evaluate thermoelectric
voltages $V_{1}$ and $V_{2}$ it is necessary to bear in mind that no currents
can flow trough the normal terminals, i.e. $I_{N_{1}}=I_{N_{2}}=0$. Making use
of these conditions, in the limit $E_{\mathrm{Th}}\ll T_{1,2}\ll\Delta$ one
arrives at the result KDZ20

$eV_{1}=-\dfrac{1}{4}\left(\dfrac{1}{T_{1}}-\dfrac{1}{T_{2}}\right)\int\dfrac{\varepsilon
d\varepsilon}{G_{N_{1}}^{n}}\Biggl{\\{}\dfrac{G_{N_{1}}^{\mathcal{Y}}G_{N_{2}}^{L}}{G_{N_{1}}^{L}+G_{N_{2}}^{L}}+\dfrac{G_{N_{2}}^{\mathcal{Y}}G_{N_{1}}^{L}-G_{N_{1}}^{\mathcal{Y}}G_{N_{2}}^{L}}{\det|\hat{G}_{S_{1}}+\hat{G}_{S_{2}}+\hat{G}_{N_{1}}+\hat{G}_{N_{2}}|}\\\
\times\Biggl{[}G_{N_{1}}^{T}-G_{N_{1}}^{n}\dfrac{G_{S_{1}}+G_{S_{2}}}{G^{n}_{S_{1}}+G^{n}_{S_{2}}}+\dfrac{G_{N_{1}}^{\mathcal{Y}}(G_{N_{1}}^{\mathcal{Y}}+G_{N_{2}}^{\mathcal{Y}})}{G_{N_{1}}^{L}+G_{N_{2}}^{L}}\Biggr{]}\Biggr{\\}},$
(19)

which allows to express $V_{1}$ in terms of the spectral conductances. The
thermoelectric voltage $V_{2}$ is determined simply by interchanging the
indices $1\leftrightarrow 2$ in Eq. (19).

We observe that the voltages $V_{1,2}$ may differ from zero only provided the
conductances $G_{N_{1,2}}^{\mathcal{Y}}$ do not vanish, which is the case in
the presence of electron-hole asymmetry. It is straightforward to demonstrate
KDZ20  that under extra symmetry conditions (i) $L_{S_{1}}=L_{S_{2}}$ and (ii)
${\mathcal{A}}_{S_{1}}={\mathcal{A}}_{S_{2}}$ the kinetic coefficient
${\mathcal{Y}}$ equals to zero everywhere in the N-wires attached to normal
terminals N1 and N2. Hence, no electron-hole asymmetry occurs in this case and
$V_{1,2}\equiv 0$. The Josephson current across our $X$-junction was analyzed
in Ref. KDZ20  only under the conditions (i) and (ii), i.e. in the absence of
any induced thermoelectric potentials.

In this work we lift both symmetry conditions (i) and (ii) and evaluate the
supercurrent $I_{S}$ in the presence of electron-hole asymmetry and the
thermoelectric effect. The magnitude of thermoelectric potentials $|V_{1,2}|$
may be not small in this case, in some special limits reaching the values of
up to $\sim E_{\mathrm{Th}}/e$ KDZ20 . Hence, this effect should in general be
included into our consideration.

Employing the quasiclassical formalism outlined in the previous sections,
after some algebra we obtain

$I_{S}=-\dfrac{1}{4e}\int\dfrac{G_{S_{1}}-G_{S_{2}}}{\det|\hat{G}_{S_{1}}+\hat{G}_{S_{2}}+\hat{G}_{N_{1}}+\hat{G}_{N_{2}}|}\Bigl{\\{}[(G^{L}_{N_{1}}+G^{L}_{N_{2}})G^{T}_{N_{1}}+(G^{\mathcal{Y}}_{N_{1}}+G^{\mathcal{Y}}_{N_{2}})G^{\mathcal{Y}}_{N_{1}}]h^{T}_{N_{1}}\\\
-[G^{L}_{N_{1}}G^{\mathcal{Y}}_{N_{2}}-G^{\mathcal{Y}}_{N_{1}}G^{L}_{N_{2}}](h^{L}_{N_{1}}-h^{L}_{N_{2}})+[(G^{L}_{N_{1}}+G^{L}_{N_{2}})G^{T}_{N_{2}}+(G^{\mathcal{Y}}_{N_{1}}+G^{\mathcal{Y}}_{N_{2}})G^{\mathcal{Y}}_{N_{2}}]h^{T}_{N_{2}}\Bigr{\\}}d\varepsilon\\\
-\dfrac{1}{4}\int\dfrac{2\mathcal{G}}{\det|\hat{G}_{S_{1}}+\hat{G}_{S_{2}}+\hat{G}_{N_{1}}+\hat{G}_{N_{2}}|}\Bigl{\\{}[G^{\mathcal{Y}}_{N_{2}}G^{T}_{N_{1}}-G^{\mathcal{Y}}_{N_{1}}(G^{T}_{N_{2}}+G_{S_{1}}+G_{S_{2}})]h^{T}_{N_{1}}\\\
+[(G^{T}_{N_{1}}+G^{T}_{N_{2}}+G_{S_{1}}+G_{S_{2}})G^{L}_{N_{1}}+(G^{\mathcal{Y}}_{N_{1}}+G^{\mathcal{Y}}_{N_{2}})G^{\mathcal{Y}}_{N_{1}}]h^{L}_{N_{1}}+[G^{\mathcal{Y}}_{N_{1}}G^{T}_{N_{2}}-G^{\mathcal{Y}}_{N_{2}}(G^{T}_{N_{1}}+G_{S_{1}}+G_{S_{2}})]h^{T}_{N_{2}}\\\
+[(G^{T}_{N_{1}}+G^{T}_{N_{2}}+G_{S_{1}}+G_{S_{2}})G^{L}_{N_{2}}+(G^{\mathcal{Y}}_{N_{1}}+G^{\mathcal{Y}}_{N_{2}})G^{\mathcal{Y}}_{N_{2}}]h^{L}_{N_{2}}\Bigr{\\}}d\varepsilon.$
(20)

This cumbersome expression can be cast to the form

$I_{S}=r_{N_{2}}^{n}I_{J}(T_{1},\chi)+r_{N_{1}}^{n}I_{J}(T_{2},\chi)+I_{S}^{\mathrm{ne}},$
(21)

where $r_{N_{1,2}}=R^{n}_{N_{1,2}}/(R^{n}_{N_{1}}+R^{n}_{N_{1}})$ and

$I_{J}(T,\chi)=-\dfrac{1}{2e}\int\mathcal{G}(\varepsilon)\tanh\dfrac{\varepsilon}{2T}d\varepsilon$
(22)

is the equilibrium Josephson current across our $X$-junction. At sufficiently
high temperatures $T\gg E_{\mathrm{Th}}$ Eq. (22) yields

$I_{J}=\dfrac{16\varkappa}{3+2\sqrt{2}}\dfrac{E_{\mathrm{Th}}}{eR_{n}^{a}}\left(\dfrac{2\pi
T}{E_{\mathrm{Th}}}\right)^{3/2}e^{-\sqrt{2\pi T/E_{\mathrm{Th}}}}\sin\chi,$
(23)

where we defined
$\varkappa=4\sqrt{\mathcal{A}_{S_{1}}\mathcal{A}_{S_{2}}}/(\mathcal{A}_{S_{1}}+\mathcal{A}_{S_{2}}+\mathcal{A}_{N_{1}}+\mathcal{A}_{N_{2}})$,
$R_{n}^{a}=L_{S}/(\sigma\sqrt{\mathcal{A}_{S_{1}}\mathcal{A}_{S_{2}}})$ and
$L_{S}=L_{S_{1}}+L_{S_{2}}$. For equal cross sections
${\mathcal{A}}_{S_{1}}={\mathcal{A}}_{S_{2}}$ the parameter $R_{n}^{a}$
coincides with the normal state resistance of our junction $R_{n}$ and the
result (23) immediately reduces to that of Ref. KDZ20 , cf. also Ref. ZZh .

The last term $I_{S}^{\mathrm{ne}}$ in Eq. (21) represents an extra non-
equilibrium contribution to the supercurrent which differs from zero provided
$T_{1}\neq T_{2}$. In the limit $T_{1,2}\gg E_{\mathrm{Th}}$ this term reduces
to

$I_{S}^{\mathrm{ne}}=\dfrac{1}{e}\left(\dfrac{1}{T_{1}}-\dfrac{1}{T_{2}}\right)K(\chi),$
(24)

where

$K(\chi)=-\dfrac{1}{4}\int\mathcal{G}\dfrac{(G^{L}_{N_{1}}G^{n}_{N_{2}}-G^{L}_{N_{2}}G^{n}_{N_{1}})}{(G^{L}_{N_{1}}+G^{L}_{N_{2}})(G^{n}_{N_{1}}+G^{n}_{N_{2}})}\varepsilon
d\varepsilon\\\
+\dfrac{1}{4}\int\Biggl{[}\dfrac{G_{S_{1}}G_{S_{2}}^{n}-G_{S_{2}}G_{S_{1}}^{n}}{G^{n}_{S_{1}}+G^{n}_{S_{2}}}+\mathcal{G}\dfrac{G^{\mathcal{Y}}_{N_{1}}+G^{\mathcal{Y}}_{N_{2}}}{G^{L}_{N_{1}}+G^{L}_{N_{2}}}\Biggr{]}\dfrac{[G^{L}_{N_{1}}G^{\mathcal{Y}}_{N_{2}}-G^{\mathcal{Y}}_{N_{1}}G^{L}_{N_{2}}]\varepsilon
d\varepsilon}{\det|\hat{G}_{S_{1}}+\hat{G}_{S_{2}}+\hat{G}_{N_{1}}+\hat{G}_{N_{2}}|}.$
(25)

Figure 2: Universal functions $P_{1,2}(L_{S_{1}}/L_{S})$ (left panel) and
$P_{3,4}(L_{S_{1}}/L_{S})$ (right panel). We observe that
$P_{1}(1/2)=P_{2}(1/2)=1$, while $P_{3}(1/2)=P_{4}(1/2)=-1$.

Our detailed numerical analysis of Eq. (25) demonstrates that the effect of
electron-hole asymmetry on the function $K(\chi)$, though exists, always
remains very small and, hence, can be safely neglected for any geometry of our
$X$-junction. This is an important conclusion which allows to drop the
conductance $G^{\mathcal{Y}}$ from Eq. (25) and significantly simplify our
further calculations. Setting $G^{\mathcal{Y}}$ to zero, from Eq. (25) we
obtain

$K(\chi)=-\dfrac{1}{4}\int\mathcal{G}(\varepsilon)W(\varepsilon)\varepsilon
d\varepsilon,\quad
W(\varepsilon)=\dfrac{(G^{L}_{N_{1}}G^{n}_{N_{2}}-G^{L}_{N_{2}}G^{n}_{N_{1}})}{(G^{n}_{N_{1}}+G^{n}_{N_{2}})(G^{L}_{N_{1}}+G^{L}_{N_{2}})},$
(26)

Proceeding along the lines with Ref. KDZ20  one can directly evaluate the
functions $\mathcal{G}$ and $W$ in the high energy limit. Extrapolating this
high energy expansion to the whole energy interval, evaluating the integral in
Eq. (26) and combining the result with Eq. (24), we get

$I_{S}^{\mathrm{ne}}=\dfrac{4\varkappa^{3}}{(3+2\sqrt{2})^{2}}\dfrac{1101}{1250}r_{N_{1}}r_{N_{2}}\left(\dfrac{1}{T_{1}}-\dfrac{1}{T_{2}}\right)\left(\dfrac{L_{S}}{L_{N_{2}}}-\dfrac{L_{S}}{L_{N_{1}}}\right)\dfrac{E_{\mathrm{Th}}}{eR_{n}^{a}}\\\
\times\left\\{\dfrac{\mathcal{A}_{S_{1}}}{2\mathcal{A}_{S_{2}}}P_{1}(L_{S_{1}}/L)+\dfrac{\mathcal{A}_{S_{2}}}{2\mathcal{A}_{S_{1}}}P_{1}(L_{S_{2}}/L)+P_{2}(L_{S_{1}}/L)\cos\chi\right\\}\sin\chi,$
(27)

where the universal functions $P_{1}$ and $P_{2}$ are displayed in Fig. 2
(left panel).

Equation (27) demonstrates that in a wide temperature interval
$T_{1,2}>E_{\mathrm{Th}}$ the non-equilibrium contribution to the Josephson
current is described by a universal power law dependence $\propto
1/T_{1}-1/T_{2}$, whereas the phase dependence of $I_{S}^{\mathrm{ne}}$
depends on the junction geometry only. In Fig. 3 we display the results of our
numerical solution of the Usadel equation together with those defined by Eq.
(27).

Figure 3: Phase dependence of the non-equilibrium term
$I_{S}^{\mathrm{ne}}(\chi)$, normalized by the factor
$(1/T_{1}-1/T_{2})E^{2}_{\mathrm{Th}}/eR_{n}^{a}$. Thick black lines represent
our analytical results, whereas numerically exact results are shown by thin
color lines. Top left panel: The parameters $L_{S_{1}}=0.4L_{S}$ and
$T_{1}=65E_{\mathrm{Th}}$ apply to all curves, different thin color curves
correspond to different temperatures $T_{2}$ ($40E_{\mathrm{Th}}$,
$50E_{\mathrm{Th}}$, $60E_{\mathrm{Th}}$, $70E_{\mathrm{Th}}$,
$80E_{\mathrm{Th}}$, $90E_{\mathrm{Th}}$). Top right panel: The parameters
$L_{S_{1}}=0.3L_{S}$ and $T_{1}=60E_{\mathrm{Th}}$ apply to all curves,
different thin color curves correspond to different temperatures $T_{2}$
($35E_{\mathrm{Th}}$, $45E_{\mathrm{Th}}$, $55E_{\mathrm{Th}}$,
$65E_{\mathrm{Th}}$, $75E_{\mathrm{Th}}$, $85E_{\mathrm{Th}}$). Bottom left
panel: The parameters $L_{S_{1}}=0.2L_{S}$ and$T_{1}=55E_{\mathrm{Th}}$ apply
to all curves, different thin color curves correspond to different
temperatures $T_{2}$ ($30E_{\mathrm{Th}}$, $40E_{\mathrm{Th}}$,
$50E_{\mathrm{Th}}$, $60E_{\mathrm{Th}}$, $70E_{\mathrm{Th}}$,
$80E_{\mathrm{Th}}$). Bottom right panel: The parameters $L_{S_{1}}=0.1L_{S}$
and $T_{1}=50E_{\mathrm{Th}}$ apply to all curves, different thin color curves
correspond to different temperatures $T_{2}$ ($25E_{\mathrm{Th}}$,
$35E_{\mathrm{Th}}$, $45E_{\mathrm{Th}}$, $55E_{\mathrm{Th}}$,
$65E_{\mathrm{Th}}$, $75E_{\mathrm{Th}}$). Other parameters are:
$\Delta=1000E_{\mathrm{Th}}$, $L_{N_{1}}=L_{S}$, $L_{N_{2}}=3L_{S}$,
$\mathcal{A}_{S_{1}}=\mathcal{A}_{S_{2}}=\mathcal{A}_{N_{1}}=\mathcal{A}_{N_{2}}$.

For partially symmetric junction with $L_{S_{1}}=L_{S_{2}}$ and
$\mathcal{A}_{S_{1}}=\mathcal{A}_{S_{2}}$ we have $P_{1}=P_{2}=1$ and the
phase dependence of $I_{S}^{\mathrm{ne}}$ reduces to the peculiar form
$\cos^{2}(\chi/2)\sin\chi$ KDZ20 . For strongly asymmetric junctions the
function $P_{1}$ increases by about an order of magnitude, whereas $P_{2}$
varies slightly (see Fig. 2). In this case the current-phase relation
approaches a sin-like form. The results displayed in Fig. 3 demonstrate that
our analytic formula (27) is in a good agreement with numerically exact
results for $I_{S}^{\mathrm{ne}}$ as long as the lengths $L_{S_{1,2}}$ remain
not very small $\min(L_{S_{1}},L_{S_{2}})\gtrsim 0.2L_{S}$.

##  5 Supercurrent stimulation and $\pi$-junction states

The above results demonstrate that in the presence of a temperature gradient
and at sufficiently high temperatures the supercurrent is not anymore
exponentially small due to the non-equilibrium contribution (27). In other
words, at $T_{1,2}\gg E_{\mathrm{Th}}$ the Josephson current in our
$X$-junction is stimulated by the temperature gradient. The relative magnitude
of $I_{S}^{\mathrm{ne}}$ – as compared to the equilibrium term $I_{J}$ (23) –
can further be enhanced by a proper choice of the junction geometry. For
instance, in a strongly asymmetric limiting case $L_{S_{1}}\ll L_{S_{2}}$ and
$\mathcal{A}_{S_{1}}\gg\mathcal{A}_{S_{2}}+\mathcal{A}_{N_{1}}+\mathcal{A}_{N_{2}}$
we obtain

$I_{S}^{\mathrm{ne}}=\dfrac{32\varkappa}{(3+2\sqrt{2})^{2}}\dfrac{1101}{1250}r_{N_{1}}r_{N_{2}}\left(\dfrac{1}{T_{1}}-\dfrac{1}{T_{2}}\right)\left(\dfrac{L_{S}}{L_{N_{2}}}-\dfrac{L_{S}}{L_{N_{1}}}\right)\dfrac{E_{\mathrm{Th}}}{eR_{n}^{a}}P_{1}(L_{S_{1}}/L)\sin\chi.$
(28)

It is easy to verify that this term dominates the supercurrent already at
$\min T_{1,2}\gtrsim 30E_{\mathrm{Th}}$, i.e. at considerably lower
temperatures than in symmetric junctions with $L_{S_{1}}=L_{S_{2}}$ and
$\mathcal{A}_{S_{1}}=\mathcal{A}_{S_{2}}=\mathcal{A}_{N_{1}}=\mathcal{A}_{N_{2}}$,
where the analogous condition reads $T_{1,2}\gtrsim 70E_{\mathrm{Th}}$, cf.
Ref. KDZ20 . We also observe from Figs. 4 and 5 below that for partially
asymmetric junctions with $L_{S_{1}}\ll L_{S_{2}}$ and
$\mathcal{A}_{S_{1}}=\mathcal{A}_{S_{2}}=\mathcal{A}_{N_{1}}=\mathcal{A}_{N_{2}}$
nonequilibrium effects become visible only at $T_{1,2}/E_{\mathrm{Th}}\gtrsim
40\div 50$. Thus, we conclude that strongly asymmetric $X$-junctions with
$L_{S_{1}}\ll L_{S_{2}}$ and
$\mathcal{A}_{S_{1}}\gg\mathcal{A}_{S_{2}}+\mathcal{A}_{N_{1}}+\mathcal{A}_{N_{2}}$
appear to be most suitable candidates for observing the non-equilibrium
Josephson current defined in Eq. (28).

Our general analysis also applies provided the temperature difference becomes
large. In order to address this particular situation let us set $T_{1}\to 0$
and $T_{2}\gg E_{\mathrm{Th}}$. As before, the effect of electron-hole
asymmetry on $I_{S}$ remains weak and can be neglected in this case. Then Eq.
(21) reduces to the form (cf. also Ref. KDZ20 ):

$I_{S}=r_{N_{2}}^{n}I_{J}(0,\chi)+I^{\mathrm{ne}}_{S},\quad
I^{\mathrm{ne}}_{S}=-\dfrac{1}{2e}\int\mathcal{G}(\varepsilon)W(\varepsilon)\operatorname{sgn}\varepsilon
d\varepsilon.$ (29)

In order to obtain a simple analytic estimate for the equilibrium Josephson
current at zero temperature $I_{J}(0,\chi)$ we again employ the high energy
solution of the Usadel equation extending it to all energies. Then from Eq.
(22) we get

$I_{J}(0,\chi)=\dfrac{32\varkappa}{3+2\sqrt{2}}\dfrac{E_{\mathrm{Th}}}{eR_{n}^{a}}\sin\chi.$
(30)

Remarkably, in the case of SNS junctions (with $\varkappa=2$ and
$R_{n}^{a}\equiv R_{n}$) this simple estimate provides a very accurate value
of the critical current $I_{C}(0)$ (cf., Ref. GreKa ). As expected, it cannot
capture the non-sinusoidal character of the current-phase relation in this
case, which is, however, unimportant for our present purposes.

Evaluating the non-equilibrium term $I^{\mathrm{ne}}_{S}$ we employ the same
approximation combined with the expression for $W(\varepsilon)$ (26) and find

$I^{\mathrm{ne}}_{S}=\dfrac{0.72\varkappa^{3}}{(3+2\sqrt{2})^{2}}r_{N_{1}}r_{N_{2}}\left(\dfrac{L_{S}}{L_{N_{2}}}-\dfrac{L_{S}}{L_{N_{1}}}\right)\dfrac{E_{\mathrm{Th}}}{eR_{n}^{a}}\Biggl{\\{}\dfrac{\mathcal{A}_{S_{1}}}{\mathcal{A}_{S_{2}}}P_{3}(L_{S_{1}}/L)+\\\
+\dfrac{\mathcal{A}_{S_{2}}}{\mathcal{A}_{S_{1}}}P_{3}(L_{S_{2}}/L)+2P_{4}(L_{S_{1}}/L)\cos\chi\Biggr{\\}}\sin\chi,$
(31)

where the functions $P_{3}$ and $P_{4}$ are displayed in Fig. 2 (right panel).

In particular, in strongly asymmetric $X$-junctions with $L_{S_{1}}\ll
L_{S_{2}}$ and
$\mathcal{A}_{S_{1}}\gg\mathcal{A}_{S_{2}}+\mathcal{A}_{N_{1}}+\mathcal{A}_{N_{2}}$
this result reduces to

$I^{\mathrm{ne}}_{S}=\dfrac{11.52\varkappa}{(3+2\sqrt{2})^{2}}r_{N_{1}}r_{N_{2}}\left(\dfrac{L_{S}}{L_{N_{2}}}-\dfrac{L_{S}}{L_{N_{1}}}\right)\dfrac{E_{\mathrm{Th}}}{eR_{n}^{a}}P_{3}(L_{S_{1}}/L)\sin\chi.$
(32)

We observe that for large temperature gradients the magnitude of the non-
equilibrium contribution to the Josephson current $I_{S}^{\mathrm{ne}}$ –
apart from some geometry factors – can be of the same order as that of the
equilibrium supercurrent $I_{J}$ (30). These two contributions to $I_{S}$ can
have either the same or opposite signs depending on whether $L_{N_{1}}$ is
longer or shorter than $L_{N_{2}}$.

Figure 4: Left panel: Josephson critical current $I_{C}\equiv\max|I_{S}|$ as a
function of $T_{2}$. Solid lines correspond to our numerically exact solution,
dashed lines indicate the result (21) combined with (24) and (27), dotted line
represent the quasi-equilibrium contribution
$r_{N_{2}}I_{J}(T_{1},\pi/2)+r_{N_{1}}I_{J}(T_{2},\pi/2)$ to $I_{S}$. Right
panel: The phase dependencies of the Josephson current $I_{S}(\chi)$ for
$T_{2}=50E_{\mathrm{Th}}$ (“a” curves) and $T_{2}=70E_{\mathrm{Th}}$ (“b”
curves). Other parameters are the same for both panels:
$T_{1}=55E_{\mathrm{Th}}$, $L_{N_{1}}=L_{S}$, $L_{N_{2}}=3L_{S}$,
$L_{S_{1}}=0.2L_{S}$ and
$\mathcal{A}_{S_{1}}=\mathcal{A}_{S_{2}}=\mathcal{A}_{N_{1}}=\mathcal{A}_{N_{2}}$.

In order to analyze the temperature dependence of the critical Josephson
current $I_{C}=\max I_{S}(\chi)$ we will again employ Eq. (21). At the
temperatures exceeding the Thouless energy we have $I_{J}\propto\sin\chi$ with
a positive prefactor (see Eq. (23)), whereas the phase dependence of the non-
equilibrium term $I_{S}^{\mathrm{ne}}$ (27) is described by a somewhat
distorted $\sin\chi$ function with either positive or negative sign in front
of it depending on the sign of the product
$(T_{2}-T_{1})(L_{N_{1}}-L_{N_{2}})$.

Keeping one of the two temperatures (e.g., $T_{1}$) fixed and varying $T_{2}$
we observe that the junction behavior strongly depends on the relation between
the lengths $L_{N_{1}}$ and $L_{N_{2}}$. For $L_{N_{1}}<L_{N_{2}}$ the sign of
$I_{S}^{\mathrm{ne}}$ is remains positive at $T_{2}<T_{1}$ and turns negative
at $T_{2}>T_{1}$, implying that at sufficiently high temperatures this term
dominates over exponentially decaying contributions containing $I_{J}$ (23)
and, hence, our $X$-junction switches to the $\pi$-junction state. This
behavior is illustrated in Fig. 4 (left panel). In the right panel of Fig. 4
we display typical current-phase dependencies corresponding to both $0$\- and
$\pi$-junction states realized in our structure.

For $L_{N_{1}}>L_{N_{2}}$ our $X$-junction may already exhibit two transitions
between $0$\- and $\pi$-junction states. In this case our system remains in
the 0-junction state provided $T_{2}$ remains low enough to keep the quasi-
equilibrium term larger than $I_{S}^{\mathrm{ne}}$. However, since with
increasing $T_{2}$ (albeit still $T_{2}<T_{1}$) the contribution $\propto
I_{J}$ decays much faster as the non-equilibrium one (now having a negative
sign), the $X$-junction eventually switches to the $\pi$-junction state.
Increasing $T_{2}$ further we reach the point $T_{2}=T_{1}$ where
$I_{S}^{\mathrm{ne}}$ changes its sign, thus signaling the transition back to
the $0$-junction state at $T_{2}$ slightly below $T_{1}$. This behavior is
illustrated in Fig. 5.

Figure 5: Left panel: The same as in Fig. 4 (left panel). Right panel: The
phase dependencies of the Josephson current $I_{S}(\chi)$ for
$T_{2}=40E_{\mathrm{Th}}$ (“a” curves), $T_{2}=50E_{\mathrm{Th}}$ (“b”
curves), and $T_{2}=70E_{\mathrm{Th}}$ (“c” curves). Other parameters are the
same for both panels: $T_{1}=65E_{\mathrm{Th}}$, $L_{N_{1}}=3L_{S}$,
$L_{N_{2}}=L_{S}$, $L_{S_{1}}=0.2L_{S}$ and
$\mathcal{A}_{S_{1}}=\mathcal{A}_{S_{2}}=\mathcal{A}_{N_{1}}=\mathcal{A}_{N_{2}}$.

In summary, we evaluated the Josephson current $I_{S}$ across an $X$-junction
exposed to a temperature gradient and demonstrated that the effect of
supercurrent stimulation can be further enhanced by a proper choice of
geometric parameters for our device. We also predicted a non-trivial power-law
temperature dependence of $I_{S}$ and showed that our $X$-junction may exhibit
transitions between 0- and $\pi$-junction states controlled both by the
temperature gradient and the circuit topology. Note that the junction topology
can further be modified to embrace, e.g., the one addressed in Ref. VH  in the
context of the thermoelectric effect. The whole analysis (to be published
elsewhere KZ20 ) becomes much more involved but some of our key observations
remain applicable also in this case. It would be interesting to test our
predictions in future experiments.

This work was supported in part by RFBR grant No. 18-02-00586.

## References

  * (1) B.D. Josephson, Phys. Lett. 1, 251 (1962)
  * (2) A.D. Zaikin, G.F. Zharkov, Fiz. Nizk. Temp. 7, 375 (1981) [Sov. J. Low Temp. Phys. 7, 181 (1981)]
  * (3) P. Dubos, H. Courtois, B. Pannetier, F.K. Wilhelm, A.D. Zaikin, G. Schön, Phys. Rev. B 63, 064502 (2001)
  * (4) W. Belzig, F.K. Wilhelm, C. Bruder, G. Schön, A.D. Zaikin, Superlatt. Microstruct. 25, 1251 (1999)
  * (5) A.A. Golubov, M.Yu. Kupriyanov, E. Il’ichev, Rev. Mod. Phys. 76, 411 (2004)
  * (6) L.G. Aslamazov, S.V. Lempitskii, Zh. Eksp. Teor. Fiz. 82, 1671 (1982) [Sov. Phys. JETP 55, 967 (1982)]
  * (7) A.D. Zaikin, Zh. Eksp. Teor. Fiz. 84, 1560 (1983) [Sov. Phys. JETP 57, 910 (1983)]
  * (8) A.F. Volkov, Phys. Rev. Lett. 74, 4730 (1995)
  * (9) F.K. Wilhelm, G. Schön, A.D. Zaikin, Phys. Rev. Lett. 81, 1682 (1998)
  * (10) S. Yip, Phys. Rev. B 58, 5803 (1998)
  * (11) J.J.A. Baselmans, A.F. Morpurgo, B.J. van Wees, T.M. Klapwijk, Nature 397, 43 (1999)
  * (12) M.S. Kalenkov, P.E. Dolgirev, A.D. Zaikin, Phys. Rev. B 101, 180505(R) (2020)
  * (13) H. Courtois, P. Gandit, D. Mailly, B. Pannetier, Phys. Rev. Lett. 76, 130 (1996)
  * (14) A.A. Golubov, F.K. Wilhelm, A.D. Zaikin, Phys. Rev. B 55, 1123 (1997)
  * (15) V.L. Ginzburg, Rev. Mod. Phys. 76, 981 (2004)
  * (16) P. Virtanen, T.T. Heikkilä, Phys. Rev. Lett. 92, 177004 (2004)
  * (17) M.S. Kalenkov, A.D. Zaikin, in preparation


# Representation Matters: Offline Pretraining
for Sequential Decision Making

Mengjiao Yang Ofir Nachum<EMAIL_ADDRESS><EMAIL_ADDRESS>Google
Research Google Research

###### Abstract

The recent success of supervised learning methods on ever larger offline
datasets has spurred interest in the reinforcement learning (RL) field to
investigate whether the same paradigms can be translated to RL algorithms.
This research area, known as _offline RL_ , has largely focused on offline
policy optimization, aiming to find a return-maximizing policy exclusively
from offline data. In this paper, we consider a slightly different approach to
incorporating offline data into sequential decision-making. We aim to answer
the question, what unsupervised objectives applied to offline datasets are
able to learn state representations which elevate performance on downstream
tasks, whether those downstream tasks be online RL, imitation learning from
expert demonstrations, or even offline policy optimization based on the same
offline dataset? Through a variety of experiments utilizing standard offline
RL datasets, we find that the use of pretraining with unsupervised learning
objectives can dramatically improve the performance of policy learning
algorithms that otherwise yield mediocre performance on their own. Extensive
ablations further provide insights into what components of these unsupervised
objectives – e.g., reward prediction, continuous or discrete representations,
pretraining or finetuning – are most important and in which settings.

##  1 Introduction

Within the reinforcement learning (RL) research field, _offline RL_ has
recently gained a significant amount of interest (Levine et al., 2020; Lange
et al., 2012). Offline RL considers the problem of performing reinforcement
learning – i.e., learning a policy to solve a sequential decision-making task
– exclusively from a static, offline dataset of experience. The recent
interest in offline RL is partly motivated by the success of _data-driven_
methods in the supervised learning literature. Indeed, the last decade has
witnessed ever more impressive models learned from ever larger static datasets
(Halevy et al., 2009; Krizhevsky et al., 2012; Brown et al., 2020; Dosovitskiy
et al., 2020). Solving offline RL is therefore seen as a stepping stone
towards developing scalable, data-driven methods for policy learning (Fu et
al., 2020). Accordingly, much of the recent offline RL research focuses on
proposing new policy optimization algorithms amenable to learning from offline
datasets (e.g., Fujimoto et al. (2019); Wu et al. (2019); Agarwal et al.
(2020); Kumar et al. (2020); Yu et al. (2020); Matsushima et al. (2020)).

In this paper, we consider a slightly different approach to incorporating
offline data into sequential decision-making. We are inspired by recent
successes in semi-supervised learning (Mikolov et al., 2013; Devlin et al.,
2018; Chen et al., 2020), in which large and potentially unlabelled offline
datasets are used to learn _representations_ of the data – i.e., a mapping of
input to a fixed-length vector embedding – and these representations are then
used to accelerate learning on a downstream supervised learning task. We
therefore consider whether the same paradigm can apply to RL. Can offline
experience datasets be used to learn representations of the data that
accelerate learning on a downstream task?

Figure 1: A summary of the advantages of representation learning via
contrastive self-prediction, across a variety of settings: imitation learning,
offline RL, and online RL. Each subplot shows the aggregated mean reward and
standard error during training, with aggregation over offline datasets of
different behavior (e.g., expert, medium, etc.), with five seeds per dataset
(see Section 3). Representation learning yields significant performance gains
in all domains and tasks.

This broad and general question has been partially answered by previous works
(Ajay et al., 2020; Singh et al., 2020). These works focus on using offline
datasets to learn representations of _behaviors_ , or actions. More
specifically, these works learn a spectrum of behavior policies, conditioned
on a latent $z$, through supervised action-prediction on the offline dataset.
The latent $z$ then effectively provides an abstract action space for learning
a hierarchical policy on a downstream task, and this straightforward paradigm
is able to accelerate learning in a variety of sequential decision-making
settings. Inspired by these promising results and to differentiate our own
work, we focus our efforts on the question of representation learning for
_observations_ , or states, as opposed to learning representations of
behaviors or actions. That is, we aim to answer the question, can offline
experience datasets be used to learn representations _of state observations_
such that learning policies from these pretrained representations, as opposed
to the raw state observations, improves performance on a downstream
task?111Whether the two aspects of representation learning – action
representations and state representations – can be combined is an intriguing
question. However, to avoid an overly broad paper, we focus only on state
representation learning, and leave the question of combining this with action
representation learning to future work.

To approach this question, we devise a variety of offline datasets and
corresponding downstream tasks. For offline datasets, we leverage the Gym-
MuJoCo datasets from D4RL (Fu et al., 2020), which provide a diverse set of
datasets from continuous control simulated robotic environments. For
downstream tasks, we consider three main categories: (1) low-data imitation
learning, in which we aim to learn a task-solving policy from a small number
of expert trajectories; (2) offline RL, in which we aim to learn a task-
solving policy from the same offline dataset used for representation learning;
and (3) online RL, in which we aim to learn a task-solving policy using online
access to the environment.

Once these settings are established, we then continue to evaluate the ability
of state representation learning on the offline dataset to accelerate learning
on the downstream task. Our experiments are separated into two parts,
_breadth_ and _depth_. First for breadth, we consider a diverse variety of
representation learning objectives taken from the RL and supervised learning
literature. The results of these experiments show that, while several of these
objectives perform poorly, a few yield promising results. This promising set
essentially comprises of objectives which we call _contrastive self-
prediction_ ; these objectives take sub-trajectories of experience and then
use some components of the sub-trajectory to predict other components, with a
contrastive loss when predicting states – e.g., using a contrastive loss on
the affinity between a sequence of states and actions and the next state, akin
to popular methods in the supervised learning literature (Mikolov et al.,
2013; Devlin et al., 2018).

These initial findings guide our second set of experiments. Aiming for depth,
we devise an extensive ablation based on contrastive self-prediction to
investigate what components of the objective are most important and in which
settings. For example, whether it is important to include reward as part of
the sub-trajectory, or whether discrete representations are better than
continuous, whether pre-training and fixing the representations is better than
finetuning, etc. In short, we find that state representation learning can
yield a dramatic improvement in downstream learning. Compared to performing
policy learning from raw observations, we show that relatively simple
representation learning objectives on offline datasets can enable better and
faster learning on imitation learning, offline RL, and online RL (see Figure
1). We believe these results are especially compelling for the imitation
learning setting – where even a pretraining dataset that is far from expert
behavior yields dramatic improvement in downstream learning – and in the
offline RL setting – where we show the benefits of representation learning are
significant even when the pretraining dataset _is the same as_ the downstream
task dataset. We hope that these impressive results guide and encourage future
researchers to develop even better ways to incorporate representation learning
into sequential decision-making.

##  2 Background and Related Work

Representation learning for RL has a rich and diverse existing literature, and
we briefly review these relevant works.

#### Abstraction and Bisimulation

Traditionally, representation learning has been framed as learning or
identifying _abstractions_ of the state or action space of an environment
(Andre and Russell, 2002; Mannor et al., 2004; Dearden and Boutilier, 1997;
Abel et al., 2018). These methods aim to reduce the original environment state
and action spaces to more compact spaces by clustering those states and
actions which yield similar rewards and dynamics. Motivated by similar
intuitions, research into _bisimulation_ metrics has aimed to devise or learn
similarity functions between states (Ferns et al., 2004; Castro and Precup,
2010). While these methods originally required explicit knowledge of the
reward and dynamics functions of the environment, a number of recent works
have translated these ideas to stochastic representation learning objectives
using deep neural networks (Gelada et al., 2019; Zhang et al., 2020; Agarwal
et al., 2021). Many of these modern approaches effectively learn reward and
transition functions in the learned embedding space, and training of these
models is used to inform the learned state representations.

#### Representations in Model-Based Learning

The idea of learning latent state representations via learning reward and
dynamics models leads us to related work in the model-based RL literature.
Several recent model-based RL methods use latent state representations as a
way to simplify the model learning and policy rollout elements of model-based
policy optimization (Oh et al., 2017; Silver et al., 2018; Hafner et al.,
2020), with the rollout in latent space sometimes referred to as ‘imagination’
(Racanière et al., 2017; Hafner et al., 2019). Similar ideas have also
appeared under the label of ‘embed to control’ (Watter et al., 2015; Levine et
al., 2019). Other than learning representations through forward models, there
are also works which propose to learn _inverse_ models, in which an action is
predicted based on the representations of its preceding state and subsequent
state (Pathak et al., 2017; Shelhamer et al., 2016).

#### Contrastive Objectives

Beyond model-based representations, many previous works propose the use of
contrastive losses as a way of learning useful state representations (Wu et
al., 2018; Nachum et al., 2018; Srinivas et al., 2020; Stooke et al., 2020).
These works effectively define some notion of similarity between states and
use a contrastive loss to encourage similar states to have similar
representations. The similarity is usually based on either temporal vicinity
(pairs of states which appear in the same sub-trajectory) or user-specified
augmentations, such as random shifts of image observations (Srinivas et al.,
2020). Previous work has established connections between the use of
contrastive loss and mutual information maximization (van den Oord et al.,
2019) and energy-based models (LeCun and Huang, 2005).

#### State Representation Learning in Offline RL

The existing works mentioned above almost exclusively focus on online
settings, often learning the representations on a continuously evolving
dataset and in tandem with online policy learning. In contrast, our work
focuses on representation learning on offline datasets and separated from
downstream task learning. This serves two purposes: First, using static
offline datasets makes comparisons between different methods easier, avoiding
confounding factors associated with issues of exploration or nonstationary
datasets. Second, the offline setting is arguably more practical; in practice,
static offline datasets are more common than cheap online access to an
environment (Levine et al., 2020). Previous work in a similar vein to ours
includes Stooke et al. (2020) and Shelhamer et al. (2016), which propose to
use unsupervised pretraining, typically only on expert demonstrations, as a
way of initializing an image encoder for downstream online RL. Our own work
complements these existing studies, by presenting extensive comparisons of a
variety of representation learning objectives in several distinct settings.
Moreover, our work is unique for showing benefits of representation learning
on non-image tasks, thus avoiding the use of any explicit or implicit prior
knowledge that is typically exploited for images (e.g., using image-based
augmentations or using a convolutional network architecture).

##  3 Task Setups

We now continue to our own contributions, starting by elaborating on the
experimental protocol we design to evaluate representation learning in the
context of low-data imitation learning, offline RL (specifically, offline
policy optimization), and online RL in partially observable environments. This
protocol is summarized in Table 1.

Table 1: A summary of our experimental setups. In total, there are 16 choices
of offline data and downstream task combinations each for imitation learning,
offline RL, and online RL. Given that we run each setting with five random
seeds, this leads to a total of 240 training runs for every representation
learning objective we consider.

| Imitation |
---|---|---
Choose $\text{domain}\in\\{\text{halfcheetah},\text{hopper},\text{walker2d},\text{ant}\\}$ | $\rightarrow$ |  Offline dataset: {domain}-{data}-v0 Downstream task: Behavioral cloning (BC) on first $N$ transitions from {domain}-expert-v0
Choose $\text{data}\in\\{\text{medium},\text{medium-replay}\\}$
Choose $N\in\\{10000,25000\\}$
| Offline RL |
Choose $\text{domain}\in\\{\text{halfcheetah},\text{hopper},\text{walker2d},\text{ant}\\}$ | $\rightarrow$ |  Offline dataset: {domain}-{data}-v0 Downstream task: Behavior regularized actor critic (BRAC) on data from {domain}-{data}-v0
Choose $\text{data}\in\\{\text{expert},\text{medium-expert},\text{medium},$
$\text{medium-replay}\\}$
| Online RL |
Choose $\text{domain}\in\\{\text{halfcheetah},\text{hopper},\text{walker2d},\text{ant}\\}$ | $\rightarrow$ |  Offline dataset: {domain}-{data}-v0 with random masking Downstream task: Soft actor critic (SAC) on randomly masked version of {domain}
Choose $\text{data}\in\\{\text{expert},\text{medium-expert},\text{medium}\\}$
$\text{medium-replay}\\}$

###  3.1 Datasets

We leverage the Gym-MuJoCo datasets from D4RL (Fu et al., 2020). These
datasets are generated from running policies on the well-known MuJoCo
benchmarks of simulated locomotive agents: halfcheetah, hopper, walker2d, and
ant. Each of these four domains is associated with four datasets – expert,
medium-expert, medium, and medium-replay – corresponding to the quality of the
policies used to collect that data. Each dataset is composed of a number of
trajectories $\tau:=(s_{0},a_{0},r_{0},s_{1},a_{1},r_{1},\dots,s_{T})$. For
example, the dataset ant-expert-v0 is a dataset of trajectories generated by
expert task-solving policies on the ant domain, while the dataset halfcheetah-
medium-v0 is generated by mediocre, far from task-solving, policies.

Notably, although the underyling MuJoCo environments are Markovian, the
datasets are not necessarily Markovian, as they may be generated by multiple
distinct policies.

###  3.2 Imitation Learning in Low-Data Regime

Imitation learning (Hussein et al., 2017) seeks to match the behavior of an
agent with that of an expert. While expert demonstrations are often limited
and expensive to obtain in practice, non-expert experience data (e.g.,
generated from a mediocre agent randomly interacting with an environment) can
be much more easily accessible.

To mimic this practical scenario, we consider an experimental protocol in
which the downstream task is behavioral cloning (Pomerleau, 1991) on a small
set of expert trajectories – selected by taking either the first $10$k or
$25$k transitions from an expert dataset in D4RL, corresponding to about 10
and 25 expert trajectories, respectively. We then consider either the medium
or medium-replay datasets from the same domain for representation
learning.222To avoid issues of extrapolation when transferring learned
representations to the expert dataset, we include the small number of expert
demonstrations in the offline dataset during pretraining. Thus, this set of
experiments aims to determine whether representations learned from large
datasets of mediocre behavior can help elevate the performance of behavioral
cloning on a much smaller expert dataset.

###  3.3 Offline RL with Behavior Regularization

One of the main motivations for the introduction of the D4RL datasets was to
encourage research into fully offline reinforcement learning; i.e., whether it
is possible to learn return-maximizing policies exclusively from a static
offline dataset. Many algorithms for this setting have recently been proposed,
commonly employing some sort of _behavior regularization_ (Kumar et al., 2019;
Jaques et al., 2019; Wu et al., 2019). In its simplest form, behavior
regularization augments a vanilla actor-critic algorithm with a divergence
penalty measuring the divergence of the learned policy from the offline data,
thus compelling the learned policy to choose the same actions appearing in the
dataset.

While the actor and critic are typically trained with the raw observations as
input, with this next set of experiments, we aim to determine whether
representation learning can help in this regime as well. In this setting, the
pretraining and downstream datasets are the same, determined by a single
choice of domain (halfcheetah, hopper, walker2d, or ant) and data (expert,
medium-expert, medium, or medium-replay). For the downstream algorithm, we use
behavior regularized actor-critic (BRAC) (Wu et al., 2019), which is a simple
behavior regularized method employing a KL divergence penalty. Notably,
although the original BRAC paper uses different regularization strengths and
policy learning rates for different domains, we fix these to values which we
found to generally perform best (regularization strength of $1.0$ and policy
learning rate of $0.00003$).

Thus, this set of experiments aims to determine whether learning BRAC from
learned state representations is better (in terms of performance and less
dependence on hyperparameters) than learning BRAC from the raw states, even
when the state representations are learned using the same offline dataset.

###  3.4 Online RL in Partially Observable Environments

In this set of experiments, we aim to determine whether representations
learned from offline datasets can improve or accelerate learning in an online
domain. One of the most popular online RL algorithms is soft actor critic
(SAC) (Haarnoja et al., 2019). SAC is a well-performing algorithm on its own,
and so to increase the difficulty of the downstream task, we consider a simple
modification to make our domains partially observable: zero-masking out a
random dimension of the state observation. This modification also brings our
domains closer to practice, where partial observability due to flaky sensor
readings is common (Dulac-Arnold et al., 2019).

Accordingly, the offline dataset is determined by a choice of domain
(halfcheetah, hopper, walker2d, or ant) and data (expert, medium-expert,
medium, or medium-replay), with the same masking applied to this dataset.
Representations learned on this dataset are then applied downstream, where SAC
is trained on the online domain, with the representation module providing an
embedding of the masked observations of the environment within a learned
embedding space.

###  3.5 Evaluation

Each representation learning variant we evaluate is run with five seeds on
each of the experimental setups described above. Unless otherwise noted, a
single seed corresponds to an initial pretraining phase of $200$k steps, in
which a representation learning objective is optimized using batches of 256
sub-trajectories randomly sampled from the offline dataset. After pretraining,
the learned representation is fixed and applied to the downstream task, which
performs the appropriate training (BC, BRAC, or SAC) for $1$M steps. In this
downstream phase, every $10$k steps, we evaluate the learned policy on the
downstream domain environment by running it for $10$ episodes and computing
the average total return. We normalize this total return according to the
normalization proposed in Fu et al. (2020), such that a score of 0 roughly
corresponds to a random agent and a score of 100 to an expert agent. We
average the last 10 evaluations within the $1$M downstream training, and this
determines the final score for the run. To aggregate over multiple seeds and
task setups, we simply compute the average and standard error of this final
score.

##  4 Experiments: Breadth Study

We begin our empirical study with an initial assessment into the performance
of a broad set of representation learning ideas from the existing literature.

###  4.1 Representation Learning Objectives

We describe the algorithms we consider below. While it is infeasible for us to
extensively evaluate all previously proposed representation learning
objectives, our choice of objectives here aims to cover a diverse set of
recurring themes and ideas from previous work (see Section 2).

We use the notation

$\tau_{t:t+k}:=(s_{t},a_{t},r_{t},\dots,s_{t+k-1},a_{t+k-1},r_{t+k-1},s_{t+k})$

to denote a length-$(k+1)$ sub-trajectory of state observations, actions, and
rewards; we use $s_{t:t+k},a_{t:t+k},r_{t:t+k}$ to denote a subselection of
this trajectory based on states, actions, and rewards, respectively. We use
$\phi$ to denote the representation function; i.e., $\phi(s)$ is the
representation associated with state observation $s$, and
$\phi(s_{t:t+k}):=(\phi(s_{t}),\dots,\phi(s_{t+k}))$. All learned functions,
including $\phi$, are parameterized by neural networks. Unless otherwise
noted, $\phi$ is parameterized as a two-hidden layer fully-connected network
with 256 units per layer and output of dimension 256 (see further details in
Appendix A).

#### Inverse model

Given a sub-trajectory $\tau_{t:t+1}$, use $\phi(s_{t:t+1})$ to predict
$a_{t}$. That is, we train an auxiliary $f$ such that $f(\phi(s_{t:t+1}))$ is
a distribution over actions, and the learning objective is $-\log
P(a_{t}|f(\phi(s_{t:t+1})))$. This objective may be generalized to sequences
longer than $k+1=2$ as $-\log P(a_{t+k-1}|f(\phi(s_{t:t+k}),a_{t:t+k-1}))$.

#### Forward raw model

Given a sub-trajectory $\tau_{t:t+1}$, use $\phi(s_{t}),a_{t}$ to predict
$r_{t},s_{t+1}$. That is, we train an auxiliary $f,g$ such that
$f(\phi(s_{t}),a_{t})$ is a distribution over next states and
$g(\phi(s_{t}),a_{t})$ is a scalar reward prediction. The learning objective
is $||r_{t}-g(\phi(s_{t}),a_{t})||^{2}-\log P(s_{t+1}|f(\phi(s_{t}),a_{t}))$.
This objective may be generalized to sequences longer than $k+1=2$ as
$||r_{t}-g(\phi(s_{t:t+k-1}),a_{t:t+k-1})||^{2}-\log
P(s_{t+1}|f(\phi(s_{t:t+k-1}),a_{t:t+k}))$.

#### Forward latent model; a.k.a., DeepMDP (Gelada et al., 2019)

This is the same as the forward raw model, only that $f$ now describes a
distribution over next state representations. Thus, the log-probability with
respect to $f$ becomes $-\log P(\phi(s_{t+1})|f(\phi(s_{t}),a_{t}))$.

#### Forward energy model

This is the same as the forward raw model, only that $f$ is no longer a
distribution over raw states. Rather, $f$ maps $\phi(s_{t}),a_{t}$ to the same
embedding space as $\phi$ and the probability
$P(s_{t+1}|f(\phi(s_{t}),a_{t}))$ is defined in an energy-based way:

$\frac{\rho(s_{t+1})\exp\\{\phi(s_{t+1})^{\top}Wf(\phi(s_{t}),a_{t})\\}}{\mathbb{E}_{\rho}[\exp\\{\phi(\tilde{s})^{\top}Wf(\phi(s_{t}),a_{t})\\}]},$
(1)

where $W$ is a trainable matrix and $\rho$ is a non-trainable prior
distribution (we set $\rho$ to be the distribution of states in the offline
dataset).

#### (Momentum) temporal contrastive learning (TCL)

Given a sub-trajectory $\tau_{t:t+1}$, we apply a contrastive loss between
$\phi(s_{t}),\phi(s_{t+1})$. The objective is

$-\phi(s_{t+1})^{\top}W\phi(s_{t})+\log\mathbb{E}_{\rho}[\exp\\{\phi(\tilde{s})^{\top}W\phi(s_{t})\\}],$
(2)

where $W$ and $\rho$ are as in the forward energy model above. This objective
may be generalized to sequences longer than $k+1=2$ by having multiple terms
in the loss for $i=1,\dots,k$:

$-\phi(s_{t+i})^{\top}W_{i}\phi(s_{t})+\log\mathbb{E}_{\rho}[\exp\\{\phi(\tilde{s})^{\top}W_{i}\phi(s_{t})\\}].$
(3)

If momentum is used, we apply the contrastive loss between $f(\phi(s_{t}))$
and $\phi_{target}(s_{t+i})$, where $f$ is a learned function and
$\phi_{target}$ denotes a non-trainable version of $\phi$, with weights
corresponding to a slowly moving average of the weights of $\phi$, as in
Stooke et al. (2020); He et al. (2020).

#### Attentive Contrastive Learning (ACL)

Following the theme of contrastive losses and inspired by a number of works in
the RL (van den Oord et al., 2019) and NLP (Mikolov et al., 2013) literature
which apply such losses between tokens and contexts using an attention
mechanism, we devise a similar objective for our settings. Implementation-
wise, we borrow ideas from BERT (Devlin et al., 2018), namely we (1) take a
sub-trajectory $s_{t:t+k},a_{t:t+k},r_{t:t+k}$, (2) randomly mask a subset of
these, (3) pass the masked sequence into a transformer, and then (4) for each
masked input state, apply a contrastive loss between its representation
$\phi(s)$ and the transformer output at its sequential position. We use
$k+1=8$ in our implementation. Figure 3 provides a diagram of ACL.

#### Value prediction network (VPN)

Taken from Oh et al. (2017), this objective uses an RNN starting at
$\phi(s_{t})$ and inputting $a_{t:t+k}$ for $k$ steps to predict the $k$-step
future rewards and value functions. While the original VPN paper defines the
$(k+1)$-th value function in terms of a max over actions, we avoid this
potential extrapolation issue and simply use the $(k+1)$-th action provided in
the offline data. As we will elaborate on later, VPN bears similarities to ACL
in that it uses certain components of the input sequence (states and actions)
to predict other components (values).

#### Deep bisimulation for control

This objective is taken from Zhang et al. (2020), where the representation
function $\phi$ is learned to respect an L1 distance based on a bisimulation
similarity deduced from Bellman backups.

###  4.2 Results

Figure 2: Performance of downstream imitation learning, offline RL, and online
RL tasks under a variety of representation learning objectives. $x$-axis shows
aggregated average rewards (over five seeds) across the domains and datasets
described in Section 3. Methods that failed to converge are eliminated from
the results (see Appendix A). ACL is set to the default configuration that
favors imitation learning (see Section 5). When applicable, we also label
variants with $k+1\in\\{2,8\\}$. Methods above the dotted line are variants of
contrastive self-prediction. ACL performs well on imitation learning. VPN and
(momentum) TCL perform well on offline and online RL.

The results of these representation learning objectives are presented in
Figure 2. Representation learning, even before the extensive ablations we will
embark on in Section 5, on average improves downstream imitation learning,
offline RL, and online RL tasks by $1.5$x, $2.5$x, and $15\%$ respectively.
The objectives that appear to work best – ACL, (Momentum) TCL, VPN – fall
under a class of objectives we term _contrastive self-prediction_ , where
_self-prediction_ refers to the idea that certain components of a sub-
trajectory are predicted based on other components of the same sub-trajectory,
while _contrastive_ refers to the fact that this prediction should be
performed via a contrastive energy-based loss when the predicted component is
a state observation.

We also find that a longer sub-trajectory $k+1=8$ is generally better than a
short one $k+1=2$. The advantage here is presumably due to the non-Markovian
nature of the dataset. Even if the environment is Markovian, the use of
potentially distinct policies for data collection can lead to non-Markovian
data.

Despite these promising successes, there are a number of objectives which
perform poorly. Raw predictions of states (forward model) yields disappointing
results in these settings. Forward models of future representations – DeepMDP,
Bisimulation – also exhibit poor performance. This latter finding was
initially surprising to us, as many theoretical notions of state abstractions
are based on the principle of predictability of future state representations.
Nevertheless, even after extensive tuning of these objectives and attempts at
similar objectives (e.g., we briefly investigated incorporating ideas from
Hafner et al. (2020)), we were not able to achieve any better results. Even if
it is possible to find better architectures or hyperparameters, we believe the
difficulty in tuning these baselines makes them unattractive in comparison to
the simpler and better performing alternatives.

##  5 Experiments: Depth Study

Table 2: Factors of contrastive self-prediction considered in our ablation study and summaries of their effects. Input action and input reward default to true. The remaining factors default to false. For each effect entry, $\downarrow$ means decreased performance, $\uparrow$ means improved performance, and $=$ means no significant effect. Factor |  Description | Imitation | Offline | Online
---|---|---|---|---
reconstruct action |  Add action prediction loss based on $\phi(s)$. | $\downarrow$ | $\uparrow$ | $\uparrow$
reconstruct reward |  Add a reward prediction loss based on $\phi(s)$. | $\downarrow$ | $\uparrow$ | $\uparrow$
predict action |  Add an action prediction loss based on transformer outputs. Whenever this is true, we also set ‘input embed’ to true. | $\downarrow$ | $\uparrow$ | $\uparrow$
predict reward |  Add a reward prediction loss based on transformer outputs. Whenever this is true, we also set ‘input embed’ to true. | $\downarrow$ | $\uparrow$ | $\uparrow$
input action |  Include actions in the input sequence to transformer. | $\downarrow$ | $\uparrow$ | $\uparrow$
input reward |  Include rewards in the input sequence to transformer. | $\downarrow$ | $\uparrow$ | $\uparrow$
input embed |  Use representations $\phi(s)$ as input to transformer, as opposed to raw observations. | $\downarrow$ | $=$ | $\uparrow$
bidirectional |  To generate sequence output at position $i$, use full input sequence as opposed to only inputs at position $\leq i$. | $\downarrow$ | $=$ | $\uparrow$
finetune |  Pass gradients into $\phi$ during learning on downstream tasks. | $\downarrow$ | $\downarrow$ | $\uparrow$
auxiliary loss |  Use representation learning objective as an auxiliary loss during downstream learning, as opposed to pretraining. | $\downarrow$ | $\downarrow$ | $\uparrow$
momentum |  Adopt an additional momentum representation network. Whenever this is true, we also set ‘input embed’ to true. | $\downarrow$ | $\downarrow$ | $\uparrow$
discrete embedding |  Learn discrete representations. Following Hafner et al. (2020), we treat the 256-dim output of $\phi$ as logits to sample 16 categorical distributions of dimension 16 each and use straight-through gradients. | $\downarrow$ | $\downarrow$ | $\downarrow$
context embedding |  Following Devlin et al. (2018), use transformer output as representations for downstream tasks. Whenever this is true, we also set ‘input embed’ to true. | $\downarrow$ | $\downarrow$ | $\downarrow$
Figure 3: A pictoral representation of our depth study based on contrastive
self-prediction. We use the transformer-based architecture of attentive
contrastive learning (ACL) as a skeleton for ablations with respect to various
representation learning details. Solid arrows correspond to the configuration
of ACL. Dotted arrows and blue text are factors considered in the ablation
study. Gray blocks are masked state/action/reward entries. After the
pretraining phase, the representation network $\phi$ is reused for downstream
tasks, unless ‘context embedding’ is true, in which case the transformer is
used.  Figure 4: Ablation results on imitation learning, offline RL, and
online RL. $x$-axis shows average rewards and standard error aggregated over
either different Gym-MuJoCo datasets (imitation and offline RL) or domains
(online RL). Blue dotted lines show average rewards without pretraining. (T)
and (F) mean setting each factor to true or false (opposite from the default
configuration). Reconstructing, predicting, or inputting action or reward (row
2-7) impairs imitation performance but are important for offline and online
RL. Bidirectional transformer hurts imitation learning when downstream sample
size is small. Finetuning and auxiliary loss can help online RL. Additional
results are presented in Appendix B.

The favorable results of objectives based on the idea of contrastive self-
prediction is compelling, but the small number of objectives evaluated leaves
many questions unanswered. For example, when generating the context embedding
for a specific prediction, should one use past states (as in TCL and Momentum
TCL) or also include actions and/or rewards (as in ACL and VPN)? Should this
context use the same representation network $\phi$ (as in TCL and VPN), a
momentum version of it (as in Momentum TCL), or a completely separate network
(as in ACL)?

We use this section to study these and other important questions by conducting
a series of ablations on the factors which compose a specific contrastive
self-prediction objective and how it is applied to downstream learning. We
describe all these factors in Table 2, as well as a high-level summary of
their effects. Further anecdotal observations found during our research are
summarized in Appendix C.

We choose the transformer-based implementation of ACL to serve as the skeleton
for all these ablations (see Figure 3), due to its general favorable empirical
performance in the previous section, as well as its ease of modification. For
each downstream task below, we present the ablations with respect to the
_default_ configuration of the factors in Table 2 that corresponds to the
original ACL introduced in Section 4, and change one factor at a time to
observe its effect on downstream task performance.

###  5.1 Results

The results of our ablation studies are presented in Figure 4, and we
highlight some of the main findings below. We also take the best performing
ablation from each row (imitation, offline RL, and online RL) and plot the
performance during training in Figure 1.

Let us first consider the effects of inclusion or prediction of actions and
rewards. We notice some interesting behavior across the different downstream
modes. Namely, it appears that imitation learning is best served by focusing
only on state contrastive learning and not including or predicting actions and
rewards, whereas the offline and online RL settings appear to benefit from
these. Due to the mixed results we initially observed from including or
predicting actions and rewards, we also introduce the idea of _reconstructing_
actions and rewards based on $\phi(s)$, and we found this to have much more
consistent benefit in the RL settings, although it still degrades imitation
learning performance. This disconnect between objectives which are good for
imitation learning vs. RL, first seen in Section 4, thus continues to be
present in these ablations as well, and we find that no single objective
dominates in all settings.

We also evaluate a number of representation learning paradigms popular in the
NLP literature (Devlin et al., 2018), namely using bidirectional transformers,
finetuning, and context embedding. Although these techniques are ubiquitous in
the NLP literature, we find mixed results in RL settings. Context embedding
consistently hurts performance. Bidirectional transformer hurts imitation
learning but helps online RL. Finetuning leads to a modest degredation in
performace in imitation and offline RL but can improve online RL depending on
the domain being evaluated.

We additionally considered using the representation learning objective as an
auxiliary training loss, which is popular in the online RL literature
(Shelhamer et al., 2016; Stooke et al., 2020). And indeed, we find that it can
dramatically improve representation learning in online RL, but at the same
time, dramatically degrade performance in the offline settings (imitation
learning or offline RL).

##  6 Conclusion

Overall, our results show that relatively simple representation learning
objectives can dramatically improve downstream imitation learning, offline RL,
and online RL (Figure 1). Interestingly, our results suggest that the ideal
representation learning objective may depend on the nature of the downstream
task, and no single objective appears to dominate generally. Our extensive
ablations also provide a number of intriguing insights, showing that
representational paradigms which are popular in NLP or online RL may not
translate to good performance in offline settings.

Even with this multitude of fresh insight into the question of representation
learning in RL, our study is limited in a number of aspects, and these aspects
can serve as a starting point for future work. For example, one may consider
additional downstream tasks such as multi-task, transfer, or exploration.
Alternatively, one can extend our ablations to real-world domains like robot
learning. Or, one may consider ablating over different network architectures.

Despite these limitations, we hope our current work proves useful to RL
researchers, and serves as a guide for developing even better and more general
representation learning objectives.

## Acknowledgements

We thank Ilya Kostrikov for providing his codebase as a starting point for our
experiments as well as helpful comments on early drafts of this work. We also
thank the rest of the Google Brain team for insightful thoughts and
discussions.

## References

  * Abel et al. (2018) David Abel, Dilip Arumugam, Lucas Lehnert, and Michael Littman.  State abstractions for lifelong reinforcement learning.  In _International Conference on Machine Learning_ , pages 10–19. PMLR, 2018.
  * Agarwal et al. (2020) Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi.  An optimistic perspective on offline reinforcement learning, 2020.
  * Agarwal et al. (2021) Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Marc G Bellemare.  Contrastive behavioral similarity embeddings for generalization in reinforcement learning.  _arXiv preprint arXiv:2101.05265_ , 2021.
  * Ajay et al. (2020) Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofir Nachum.  Opal: Offline primitive discovery for accelerating offline reinforcement learning.  _arXiv preprint arXiv:2010.13611_ , 2020.
  * Andre and Russell (2002) David Andre and Stuart J Russell.  State abstraction for programmable reinforcement learning agents.  In _Aaai/iaai_ , pages 119–125, 2002.
  * Brown et al. (2020) Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _arXiv preprint arXiv:2005.14165_ , 2020.
  * Castro and Precup (2010) Pablo Castro and Doina Precup.  Using bisimulation for policy transfer in mdps.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 24, 2010.
  * Chen et al. (2020) Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pages 1597–1607. PMLR, 2020.
  * Dearden and Boutilier (1997) Richard Dearden and Craig Boutilier.  Abstraction and approximate decision-theoretic planning.  _Artificial Intelligence_ , 89(1-2):219–283, 1997\.
  * Devlin et al. (2018) 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.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Dulac-Arnold et al. (2019) Gabriel Dulac-Arnold, Daniel Mankowitz, and Todd Hester.  Challenges of real-world reinforcement learning, 2019.
  * Ferns et al. (2004) Norm Ferns, Prakash Panangaden, and Doina Precup.  Metrics for finite markov decision processes.  In _UAI_ , volume 4, pages 162–169, 2004.
  * Fu et al. (2020) Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_ , 2020.
  * Fujimoto et al. (2019) Scott Fujimoto, David Meger, and Doina Precup.  Off-policy deep reinforcement learning without exploration.  In _International Conference on Machine Learning_ , pages 2052–2062, 2019.
  * Gelada et al. (2019) Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare.  Deepmdp: Learning continuous latent space models for representation learning.  In _International Conference on Machine Learning_ , pages 2170–2179. PMLR, 2019.
  * Haarnoja et al. (2019) Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine.  Soft actor-critic algorithms and applications, 2019.
  * Hafner et al. (2019) Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.  Dream to control: Learning behaviors by latent imagination.  _arXiv preprint arXiv:1912.01603_ , 2019.
  * Hafner et al. (2020) Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba.  Mastering atari with discrete world models.  _arXiv preprint arXiv:2010.02193_ , 2020.
  * Halevy et al. (2009) Alon Halevy, Peter Norvig, and Fernando Pereira.  The unreasonable effectiveness of data.  _IEEE Intelligent Systems_ , 24(2):8–12, 2009\.
  * He et al. (2020) Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 9729–9738, 2020.
  * Hussein et al. (2017) Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne.  Imitation learning: A survey of learning methods.  _ACM Computing Surveys (CSUR)_ , 50(2):1–35, 2017\.
  * Jaques et al. (2019) Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard.  Way off-policy batch deep reinforcement learning of implicit human preferences in dialog.  _arXiv preprint arXiv:1907.00456_ , 2019.
  * Krizhevsky et al. (2012) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  _Advances in neural information processing systems_ , 25:1097–1105, 2012.
  * Kumar et al. (2019) Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine.  Stabilizing off-policy q-learning via bootstrapping error reduction.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Kumar et al. (2020) Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  _arXiv preprint arXiv:2006.04779_ , 2020.
  * Lange et al. (2012) Sascha Lange, Thomas Gabel, and Martin Riedmiller.  Batch reinforcement learning.  In _Reinforcement learning_ , pages 45–73. Springer, 2012.
  * LeCun and Huang (2005) Yann LeCun and Fu Jie Huang.  Loss functions for discriminative training of energy-based models.  In _AIStats_ , volume 6, page 34. Citeseer, 2005.
  * Levine et al. (2019) Nir Levine, Yinlam Chow, Rui Shu, Ang Li, Mohammad Ghavamzadeh, and Hung Bui.  Prediction, consistency, curvature: Representation learning for locally-linear control.  _arXiv preprint arXiv:1909.01506_ , 2019.
  * Levine et al. (2020) Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  _arXiv preprint arXiv:2005.01643_ , 2020.
  * Mannor et al. (2004) Shie Mannor, Ishai Menache, Amit Hoze, and Uri Klein.  Dynamic abstraction in reinforcement learning via clustering.  In _Proceedings of the twenty-first international conference on Machine learning_ , page 71, 2004.
  * Matsushima et al. (2020) Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu.  Deployment-efficient reinforcement learning via model-based offline optimization.  _arXiv preprint arXiv:2006.03647_ , 2020.
  * Mikolov et al. (2013) Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean.  Efficient estimation of word representations in vector space.  _arXiv preprint arXiv:1301.3781_ , 2013.
  * Nachum et al. (2018) Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine.  Near-optimal representation learning for hierarchical reinforcement learning.  _arXiv preprint arXiv:1810.01257_ , 2018.
  * Oh et al. (2017) Junhyuk Oh, Satinder Singh, and Honglak Lee.  Value prediction network.  _arXiv preprint arXiv:1707.03497_ , 2017.
  * Pathak et al. (2017) Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  In _International Conference on Machine Learning_ , pages 2778–2787. PMLR, 2017.
  * Pomerleau (1991) Dean A Pomerleau.  Efficient training of artificial neural networks for autonomous navigation.  _Neural computation_ , 3(1):88–97, 1991.
  * Racanière et al. (2017) Sébastien Racanière, Théophane Weber, David P Reichert, Lars Buesing, Arthur Guez, Danilo Rezende, Adria Puigdomenech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, et al.  Imagination-augmented agents for deep reinforcement learning.  In _Proceedings of the 31st International Conference on Neural Information Processing Systems_ , pages 5694–5705, 2017.
  * Ramachandran et al. (2017) Prajit Ramachandran, Barret Zoph, and Quoc V. Le.  Searching for activation functions, 2017.
  * Shelhamer et al. (2016) Evan Shelhamer, Parsa Mahmoudieh, Max Argus, and Trevor Darrell.  Loss is its own reward: Self-supervision for reinforcement learning.  _CoRR_ , abs/1612.07307, 2016.  URL http://arxiv.org/abs/1612.07307.
  * Silver et al. (2018) David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al.  A general reinforcement learning algorithm that masters chess, shogi, and go through self-play.  _Science_ , 362(6419):1140–1144, 2018.
  * Singh et al. (2020) Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine.  Parrot: Data-driven behavioral priors for reinforcement learning.  _arXiv preprint arXiv:2011.10024_ , 2020.
  * Srinivas et al. (2020) Aravind Srinivas, Michael Laskin, and Pieter Abbeel.  Curl: Contrastive unsupervised representations for reinforcement learning, 2020.
  * Stooke et al. (2020) Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin.  Decoupling representation learning from reinforcement learning, 2020.
  * van den Oord et al. (2019) Aaron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding, 2019.
  * Watter et al. (2015) Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller.  Embed to control: A locally linear latent dynamics model for control from raw images.  _arXiv preprint arXiv:1506.07365_ , 2015.
  * Wu et al. (2018) Yifan Wu, George Tucker, and Ofir Nachum.  The laplacian in rl: Learning representations with efficient approximations.  _arXiv preprint arXiv:1810.04586_ , 2018.
  * Wu et al. (2019) Yifan Wu, George Tucker, and Ofir Nachum.  Behavior regularized offline reinforcement learning.  _arXiv preprint arXiv:1911.11361_ , 2019.
  * Yu et al. (2020) Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma.  Mopo: Model-based offline policy optimization.  _arXiv preprint arXiv:2005.13239_ , 2020.
  * Zhang et al. (2020) Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine.  Learning invariant representations for reinforcement learning without reconstruction.  _arXiv preprint arXiv:2006.10742_ , 2020.

Appendix

##  Appendix A Experimental Details

###  A.1 Representation Network

We parametrize the representation function $\phi$ as a two-hidden layer fully
connected neural network with $256$ units per layer and output dimension
$256$. A Swish (Ramachandran et al., 2017) activation function is applied to
the output of each hidden layer. We experimented with representation dimension
sizes $16,64,256,$ and $512$, and found $256$ and $512$ to generally work the
best (see Figure 5 in Appendix B.1).

###  A.2 Transformer Network

The BERT-style transformer used in attentive contrastive learning (ACL)
consists of one preprocessing layer of $256$ units and ReLU activaiton,
followed by a multi-headed attention layer ($4$ heads with $128$ units each),
followed by a fully connected feed forward layer with hidden dimension $256$
and ReLU activation, finally followed by an output layer of $256$ units (the
same as $\phi$’s output). We experimented with different number of attention
blocks and number of heads in each block, but did not observe significant
difference in performance.

When masking input (sequences of state, actions, or rewards), we randomly
choose to ‘drop’ each item with probability $0.3$, ’switch’ with probability
$0.15$, and ‘keep’ with probability $0.15$. ‘Drop’ refers to replacing the
item with a trainable variable of the same dimension. ‘Switch’ refers to
replacing the item with a randomly sampled item from the same batch. ‘Keep’
refers to leaving the item unchanged. These probability rates where chosen
arbitrarily and not tuned.

###  A.3 Action Prediction and Reconstruction

Whenever a loss includes action prediction or reconstruction, we follow
Haarnoja et al. (2019), and (1) utilize an output distribution given by a
tanh-squashed Gaussian and (2) apply an additive adaptive entropy regularizer
to the action prediction loss.

###  A.4 Other Networks

With few exceptions, all other functions $f,g$ mentioned in Section 4 are two-
hidden layer fully connected neural networks with $256$ units per layer and
using a Swish (Ramachandran et al., 2017) activation.

The only exception is Momentum TCL, where $f$ is the same structure but using
a residual connection on the output.

###  A.5 Training

During pretraining, we use the Adam optimizer with learning rate $0.0001$,
except for the TCL variants, for which we found $0.0003$ to work better. For
Momentum TCL, we use a moving average with rate $0.05$.

###  A.6 Convergence Failures

Representations learned under objectives including forward-raw model, VPN
(with $k+1=2$), and DeepMDP consistently diverge and output NaNs on offline
and online RL, and are therefore removed from the results in Figure 2. The
bisimulation objective on offline and online RL fails to converge in some runs
but occasionally succeeds, therefore the means of succeeded runs are computed
and shown in Figure 2.

##  Appendix B Additional Experimental Results

###  B.1 Ablation Over Representation Size

Figure 5: Average reward across domains and datasets with different
representation dimensions. $256$ and $512$ work the best (this ablation is
conducted with “reconstruct action” and “reconstruct reward” set to true).

###  B.2 Additional Contrastive Learning Results

Figure 6: Additional training curves of contrastive learning objectives
aggregated over different offline datasets in the same domain. Both in this
figure and in Figure 1, we plot the best variant of ACL according to the
ablation study, namely we set “input reward” to false in imitation learning,
“reconstruct action” to true in offline RL, and “auxiliary loss” (in ant and
halfcheetah) or “finetuning” (in hopper and walker2d) to true in online RL.
The best variant of ACL generally performs the best compared to other
contrastive learning objectives, although TCL’s performance is competitive in
offline RL.

###  B.3 Ablation Results for Individual Domains and Datasets

Figure 7: Imitation learning ablation on individual domains and datasets. The
negative impact of inputting action and reward to pretraining is more evident
in halfcheetah and walker2d. Reconstructing/predicting action/reward is
especially harmful in halfcheetah, hopper, and walker2d. There always exists
some variant of ACL that is better than without representation learning (blue
lines) in all domain-dataset combinations. Figure 8: Offline RL ablation on
individual domains and datasets. The benefit of representation learning is
more evident when expert trajectories are present (e.g., expert and medium-
expert) than when they are absent (medium and medum-replay). Reconstructing
action and reward is more important in ant and halfcheetah than in hopper and
walker2d. Figure 9: Online RL ablation on individual domains and datasets.
Auxiliary loss generally improves performance in all domains and datasets.
Finetuning improves halfcheetah, hopper, and walker2d but significantly
impairs ant.

##  Appendix C Additional Anecdotal Conclusions

  1. 1.

More ablations. Although we present our ablations as only changing one factor
at a time, we also experimented with changing multiple factors at a time. We
did not find any of these additional ablations to change the overall
conclusions.

  2. 2.

Reconstruct action. One ablation that did work surprisingly well was to only
reconstruct the action (with no other loss). This appeared to perform well on
imitation learning, but poorly on other settings.

  3. 3.

More transformers. We experimented with a different application of
transformers than ACL. Namely, we attempted to treat each dimension of the
state as a token in a sequence (as opposed to using the whole state
observation as the token). We found this to provide promising results,
although it did not convincingly improve upon the configuration of ACL. Still,
it may merit further investigation by future work.

  4. 4.

Transformer architecture. We experimented with a different number of attention
blocks or number of heads in each block, but did not observe significant
differences in performance.

  5. 5.

Normalized or regularized representations. We experimented with applying an
explicit normalization layer on the output of $\phi$ and found no benefits. We
also experimented with a stochastic representation along with a KL-divergence
regularizer to the standard normal distribution, and again found no benefits.