text:
\subsection{Evaluation}

In this section, we evaluate our method on two commonly used benchmarks: \textbf{DMControl Suite}~\citep{tassa2020dmcontrol} and \textbf{Atari 2600 Games}~\citep{bellemare2013arcade}.

\subsubsection{DMControl Suite}

\textbf{Metric:} DMControl-500k, DMControl-100k\footnote{environment steps}.

\begin{table*}[ht]
\small
\centering
\begin{tabular}{cccccc}
\toprule
500K STEP SCORES & Our CURL & CURL (Paper) & DREAMER & SAC+AE & State SAC\\
\midrule
FINGER, SPIN & $828 \pm 137$ & $\mathbf{926 \pm 45}$ & $796  \pm 183$  & $884  \pm 128$ &  $923 \pm 21$ \\
CARTPOLE, SWINGUP & $808 \pm 39$ &  $\mathbf{841 \pm 45} $&  $ 762 \pm 27 $ &  $ 735 \pm 63 $ &  $ 848 \pm 15 $\\
REACHER, EASY & $\mathbf{951 \pm 26}$ & $929 \pm 44$ & $ 793 \pm 164 $ & $ 627 \pm 58 $ & $923 \pm 24 $\\
CHEETAH, RUN  & $526 \pm 59$ & $518 \pm 28$ & $\mathbf{570 \pm 253}$ & $550 \pm 34$ & $795 \pm 30$ \\
WALKER, WALK & $892 \pm 49$ & $\mathbf{902 \pm 43}$ & $897 \pm 49$ & $847 \pm 48$ & $948 \pm 54 $ \\
BALL IN CUP, CATCH & $846 \pm 103$ & $\mathbf{959 \pm 27}$ & $879 \pm 87$ & $794 \pm 58$ & $974 \pm 33$ \\
\midrule
100K STEP SCORES & Our CURL & CURL (Paper) & DREAMER & SAC+AE & State SAC\\
\midrule
FINGER, SPIN & $749 \pm 138$ & $ \mathbf{767 \pm 56}$ &  $341 \pm 70$ &  $740 \pm 64$ & $811 \pm 46$ \\
CARTPOLE, SWINGUP & $367 \pm 74$ &  $\mathbf{582 \pm 146} $ &  $326 \pm 27$ &  $311 \pm 11$ &  $835 \pm 22$\\
REACHER, EASY & $\mathbf{576 \pm 88}$ & $ 538 \pm 233 $ &  $314 \pm 155$ &  $274 \pm 14$ & $ 746 \pm 25 $\\
CHEETAH, RUN  & $225 \pm 32$ & $ \mathbf{299 \pm 48} $ &  $ 235 \pm 137$ &  $ 267 \pm 24 $ & $616 \pm 18$ \\
WALKER, WALK & $ \mathbf{410 \pm 86} $ & $ 403 \pm 24 $ & $ 277 \pm 12$ & $394 \pm 22$ & $891 \pm 82$ \\
BALL IN CUP, CATCH & $546 \pm 210 $ & $ \mathbf{769 \pm 43} $ &  $246 \pm 174$ &  $391 \pm 82$ &  $746 \pm 91$ \\
\bottomrule
\end{tabular}
\caption{Scores achieved by SCRL (mean \& standard deviation for 5 seeds) and baselines on DMControl-500k and DMControl-100k.}\label{tab:main exp table}
\end{table*}

\subsubsection{Atari 2600 Games}

\textbf{Metric}: Atari-100K

\begin{table*}[!htb]
\centering
\tiny
\begin{tabular}{lcccccccc}
\toprule
GAME & HUMAN  & RAINBOW & SIMPLE & OTRAINBOW & EFF. RAINBOW & CURL & Our CURL & RCRL \\
\midrule
ALIEN & $7127.7$ & $318.7$ & $616.9$ & $824.7$ & $739.9$ & $558.2$&$\mathbf{ 847.4}$ \\
AMIDAR & 1719.5  & 32.5 & 88.0 & 82.8 & $\mathbf{188.6}$ & 142.1 & 168.0 \\
ASSAULT & 742.0  & 231 & 527.2 & 351.9 & 431.2 & 600.6  & $\mathbf{606.8}$\\
ASTERIX & 8503.3  & 243.6 & $\mathbf{1128.3}$ & 628.5 & 470.8 & 734.5 &629.0 \\
BANK HEIST & 753.1  & 15.55 & 34.2 & $\mathbf{182.1}$ & 51.0 & 131.6 & 176.6\\
BATTLE ZONE & 37187.5  & 2360.0 & 5184.4 & 4060.6 & 10124.6 & 14870.0& $\mathbf{21820.0}$ \\
BOXING & 12.1  & -24.8 & $\mathbf{9 . 1}$ & 2.5 & 0.2 & 1.2 &3.5 \\
BREAKOUT & 30.5  & 1.2 & $\mathbf{1 6 . 4}$ & 9.84 & 1.9 & 4.9 &3.5\\
CHOPPER COMMAND & 7387.8 & 120.0 & $\mathbf{1 2 4 6 . 9}$ & 1033.33 & 861.8 & 1058.5 & 1096.0 \\
CRAZY\_CLIMBER & 35829.4  & 2254.5 & $\mathbf{6 2 5 8 3 . 6}$ & 21327.8 & 16185.3 & 12146.5&19420.0 \\
DEMON\_ATTACK & 1971.0  & 163.6 & 208.1 & 711.8 & 508.0 & 817.6&$\mathbf{860.8}$ \\
FREEWAY & 29.6 & 0.0 & 20.3 & 25.0 & $27.9$ & 26.7&$\mathbf{28.6}$ \\
FROSTBITE & 4334.7  & 60.2 & 254.7 & 231.6 & 866.8 & 1181.3&$\mathbf{2154.8}$ \\
GOPHER & 2412.5  & 431.2 & 771.0 & $\mathbf{7 7 8 . 0}$ & 349.5 & 669.3 & 426.4 \\
HERO & 30826.4  & 487 & 2656.6 & 6458.8 & 6857.0 & 6279.3& $\mathbf{6884.9}$ \\
JAMESBOND & 302.8  & 47.4 & 125.3 & 112.3 & 301.6 & $\mathbf{4 7 1 . 0}$ &368.0\\
KANGAROO & 3035.0  & 0.0 & 323.1 & 605.4 & 779.3 & 872.5 &$\mathbf{1410.0}$ \\
KRULL & 2665.5  & 1468 & $\mathbf{4 5 3 9 . 9}$ & 3277.9 & 2851.5 & 4229.6 &2928.9\\
KUNG\_FU\_MASTER & 22736.3  & $0 .$ & $\mathbf{1 7 2 5 7 . 2}$ & 5722.2 & 14346.1 & 14307.8 & 15804.0 \\
MS\_PACMAN & 6951.6  & 67 & 1480.0 & 941.9 & 1204.1 & 1465.5&$\mathbf{1687.2}$ \\
PONG & 14.6 & -20.6 & $\mathbf{12.8}$ & 1.3 & -19.3 & -16.5 &-16.8\\
PRIVATE EYE & 69571.3  & 0 & 58.3 & 100.0 & 97.8 &218.4 &$\mathbf{293.4}$ \\
QBERT & 13455.0 & 123.46 & $\mathbf{1 2 8 8 . 8}$ & 509.3 & 1152.9 & 1042.4 & 1272.5 \\
ROAD\_RUNNER & 7845.0 & 1588.46 & 5640.6 & 2696.7 & $\mathbf{9 6 0 0 . 0}$ & 5661.0 & 8890.0 \\
SEAQUEST & 42054.7  & 131.69 & $\mathbf{6 8 3 . 3}$ & 286.92 & 354.1 & 384.5 & 444.4 \\
UP\_N\_DOWN & 11693.2  & 504.6 & 3350.3 & 2847.6 & 2877.4 & 2955.2&$\mathbf{3475.6}$ \\
\midrule
Median HNS & $100.0\%$ & $0.0\%$ & $14.4\%$ &$20.4\%$ & $16.1\%$ & $17.5\%$ & $17.99\%$ &  $16.35\%$ \\
\bottomrule
\end{tabular}
\caption{Results of our method and baselines on 26 environments from Atari 2600 Games at 100k timesteps. }
\vspace{-0.3cm}
\label{tab:atari}
\end{table*}

We need some powerful sentences here to highlight our strong performance, for example, CURL outperforms prior pixel-based methods, both model-based and model-free, on complex tasks in the DeepMind Control Suite and Atari
Games showing 1.9x and 1.2x performance gains
at the 100K environment and interaction steps
benchmarks respectively (from CURL).

\subsection{Analysis on Learned Representations}

\subsection{Ablation Study on the Choices of Negative Samples}

\subsection{\textcolor{black}{More Seeds Results of Representation Analysis}}
\textcolor{black}{
For clearness, we only show the result of representation analysis with a single seed in the main text. We add the results for multiple seeds here. The detailed description of analysis task can be found in the first paragraph in Section 5.2.
}

\begin{figure}[htbp]
    \centering

    \includegraphics[width=0.5\textwidth]{analysis-curves-0.png}
    \includegraphics[width=0.5\textwidth]{analysis-curves-1.png}
    \includegraphics[width=0.5\textwidth]{analysis-curves-2.png}
    \includegraphics[width=0.5\textwidth]{analysis-curves-3.png}
    \includegraphics[width=0.5\textwidth]{analysis-curves-4.png}
    \caption{
\textcolor{black}{
Analysis of the learned representation on \emph{Alien} with five seeds.
(a) The cosine similarity between the representations of the positive/negative state-action pair and the anchor during the training of Rainbow and RCRL.
(b) The game scores of the two algorithms during the training.}
    }
    \label{fig:similarity_appendix}
\end{figure}

\newpage
\subsection{\textcolor{black}{High Data Regime Results}}
\textcolor{black}{To empirically study how applicable our model is to higher data regimes, we run the experiments on the first five Atari games (of Table~\ref{tab:Atari}) for 1.5 millon agent interactions. We show the evaluation results of both our algorithm and the rainbow baseline in Table~\ref{tab2:Atari}. We can see that RCRL outperforms the ERainbow-sa baseline for 4 out of 5 games, which may imply that our auxiliary task has the potential to improve performance in the high-data regime. }

\begin{table*}[ht]
\small
\centering
\begin{tabular}{ccccc}
\toprule
Game & ERainbow-sa~(100k) & RCRL~(100k) & ERainbow-sa~(1.5M) & RCRL~(1.5M)\\
\midrule
Alien & $813.8$ & $854.2$ & $1721$  & $\textbf{1824}$\\
Amidar & $154.2$ &  $157.7$ &  $ 398.8 $ &  $\textbf{454.5} $ \\
Assault & $576.2$ & $569.6$ & $ 572.5 $ & $ \textbf{757.9} $\\
Asterix  & $ 697 $ & $ 799 $ & $ \textbf{1370} $ & $ 1306.7 $ \\
Bank Heist & $96$ & $107.2$ & $257.3$ & $\textbf{550.7}$ \\
\bottomrule
\end{tabular}
\caption{\textcolor{black}{Scores of RCRL and ERainbow-sa on the first five Atari games.}}\label{tab2:Atari}
\end{table*}

\subsection{Evaluation on Atari and DMControl}

\textbf{Experiments on Atari.}
Our experiments on Atari are conducted in low data regime of 100k interactions between the agent and the environment, which corresponds to two hours of human play.
We show the performance of different algorithms/baselines, including
the scores for
average human (Human),
SimPLe \citep{kaiser2019model} which is a strong model-based baseline for this regime,
original Rainbow \citep{hessel2017rainbow},
Data-Efficient Rainbow \citep[ERainbow, ][]{van2019use},
ERainbow with state-action embeddings (ERainbow-sa, cf. Figure \ref{fig:nn_structure} Left),
CURL \citep{srinivas2020curl} that is based on ERainbow,
RCRL which is based on ERainbow-sa,
and the algorithm that combines the auxiliary loss for CURL and RCRL to ERainbow-sa (RCRL+CURL).
We show the evaluation results of our algorithm and the baselines on Atari games in Table \ref{tab:Atari}.
First, we observe that using state-action embedding instead of state embedding in ERainbow does not lead to significant performance change by comparing ERainbow with ERainbow-sa.
Second,
built upon ERainbow-sa, the auxiliary task in RCRL leads to better performance compared with not only the base RL algorithm but also SimPLe and CURL in terms of the median human normalized score (HNS).
Third, we can see that RCRL further boosts the learning when combined with CURL and achieves the best performance for 7 out of 26 games, which shows that our auxiliary task can be successfully combined with other auxiliary tasks that embed different sources of information to learn the representation.

\input{Atari_table}

\begin{figure}[htbp]
    \centering
    \includegraphics[width=0.9\textwidth]{figure_dmc.png}
    \caption{
Scores achieved by RCRL and other baseline algorithms during the training for different tasks in DMControl suite.
The line and the shaded area indicate the average and the standard deviation over 5 random seeds respectively.
    }
    \label{fig:dmc}
\end{figure}

\textbf{Experiments on DMControl.}
For DMControl, we compare our algorithm with the following baselines:
Pixel SAC \citep{haarnoja2018soft} which is the base RL algorithm that receives images as the input;
SLAC \citep{lee2019stochastic} that learns a latent variable model and then updates the actor and critic based on it;
SAC+AE \citep{yarats2019improving} that uses a regularized autoencoder for reconstruction in the auxiliary task; PlaNet \citep{hafner2019learning} and Dreamer \citep{hafner2019dream} that learn a latent space world model and explicitly plan with the learned model.
We also compare with a skyline, State SAC, the receives the low-dimensional state representation instead of the image as the input.
Different from Atari games, tasks in DMControl yield dense reward.
Consequently, we split the trajectories into segments using a threshold such that the difference of returns within each segment does not exceed this threshold.
Similarly, we test the algorithms in low data regime of 500k interactions.
We show the evaluation results in Figure \ref{fig:dmc}.
We can see that our auxiliary task not only brings performance improvement over the base RL algorithm but also outperforms CURL and other state-of-the-art baseline algorithms in different tasks.
Moreover, we observe that our algorithm is more robust across runs with different seeds compared with Pixel SAC and CURL (e.g., for the task \emph{Ball in cup, Catch}).

\subsection{Analysis on the learned representation}

We analyze on the learned representation of our model to demonstrate that our auxiliary task attains a representation with better generalization, which may explain why our algorithm succeeds.
We use cosine similarity to measure the generalization from one state-action pair to another in the deep learning model.
Given two state-action pairs $x_1, x_2 \in \mathcal{X}$, cosine similarity is defined as
$\frac{\phi_\theta(x_1)^T \phi_\theta(x_2)}{||\phi_\theta(x_1)||\, ||\phi_\theta(x_2)||}$, where $\phi_\theta(\cdot)$ is the learnt embedding network.

We show the cosine similarity of the representations between positive pairs (that are sampled within the same segment and therefore likely to share similar return distributions) and negative pairs (i.e., randomly sampled state-action pairs) during the training on the game \emph{Alien} in Figure \ref{fig:similarity}.
First, we observe that when a good policy is learned, the representations of positive pairs are similar while those of negative pairs are dissimilar.
This indicates that a good representation (or the representation that supports a good policy) aggregates the state-action pairs with similar return distributions.
Then, we find that our auxiliary loss can accelerate such generalization for the representation, which makes RCRL learn faster.

\begin{figure}[htbp]
    \centering
    \includegraphics[width=0.65\textwidth]{figure5.png}
    \caption{
Analysis of the learned representation on \emph{Alien}.
(a) The cosine similarity between the representations of the positive/negative state-action pair and the anchor during the training of Rainbow and RCRL.
(b) The scores of the two algorithms during the training.
    }
    \label{fig:similarity}
\end{figure}

\section{Submission of conference papers to ICLR 2021}

ICLR requires electronic submissions, processed by
\url{https://openreview.net/}. See ICLR's website for more instructions.

If your paper is ultimately accepted, the statement {\tt
  {\textbackslash}iclrfinalcopy} should be inserted to adjust the
format to the camera ready requirements.

The format for the submissions is a variant of the NeurIPS format.
Please read carefully the instructions below, and follow them
faithfully.

\subsection{Style}

Papers to be submitted to ICLR 2021 must be prepared according to the
instructions presented here.

Authors are required to use the ICLR \LaTeX{} style files obtainable at the
ICLR website. Please make sure you use the current files and
not previous versions. Tweaking the style files may be grounds for rejection.

\subsection{Retrieval of style files}

The style files for ICLR and other conference information are available online at:
\begin{center}
  \url{http://www.iclr.cc/}
\end{center}
The file \verb+iclr2021_conference.pdf+ contains these
instructions and illustrates the
various formatting requirements your ICLR paper must satisfy.
Submissions must be made using \LaTeX{} and the style files
\verb+iclr2021_conference.sty+ and \verb+iclr2021_conference.bst+ (to be used with \LaTeX{}2e). The file
\verb+iclr2021_conference.tex+ may be used as a ``shell'' for writing your paper. All you
have to do is replace the author, title, abstract, and text of the paper with
your own.

The formatting instructions contained in these style files are summarized in
sections \ref{gen_inst}, \ref{headings}, and \ref{others} below.

\section{General formatting instructions}
\label{gen_inst}

The text must be confined within a rectangle 5.5~inches (33~picas) wide and
9~inches (54~picas) long. The left margin is 1.5~inch (9~picas).
Use 10~point type with a vertical spacing of 11~points. Times New Roman is the
preferred typeface throughout. Paragraphs are separated by 1/2~line space,
with no indentation.

Paper title is 17~point, in small caps and left-aligned.
All pages should start at 1~inch (6~picas) from the top of the page.

Authors' names are
set in boldface, and each name is placed above its corresponding
address. The lead author's name is to be listed first, and
the co-authors' names are set to follow. Authors sharing the
same address can be on the same line.

Please pay special attention to the instructions in section \ref{others}
regarding figures, tables, acknowledgments, and references.

There will be a strict upper limit of 8 pages for the main text of the initial submission, with unlimited additional pages for citations. Note that the upper page limit differs from last year!Authors may use as many pages of appendices (after the bibliography) as they wish, but reviewers are not required to read these. During the rebuttal phase and for the camera ready version, authors are allowed one additional page for the main text, for a strict upper limit of 9 pages.

\section{Headings: first level}
\label{headings}

First level headings are in small caps,
flush left and in point size 12. One line space before the first level
heading and 1/2~line space after the first level heading.

\subsection{Headings: second level}

Second level headings are in small caps,
flush left and in point size 10. One line space before the second level
heading and 1/2~line space after the second level heading.

\subsubsection{Headings: third level}

Third level headings are in small caps,
flush left and in point size 10. One line space before the third level
heading and 1/2~line space after the third level heading.

\section{Citations, figures, tables, references}
\label{others}

These instructions apply to everyone, regardless of the formatter being used.

\subsection{Citations within the text}

Citations within the text should be based on the \texttt{natbib} package
and include the authors' last names and year (with the ``et~al.'' construct
for more than two authors). When the authors or the publication are
included in the sentence, the citation should not be in parenthesis using \verb|\citet{}| (as
in ``See \citet{Hinton06} for more information.''). Otherwise, the citation
should be in parenthesis using \verb|\citep{}| (as in ``Deep learning shows promise to make progress
towards AI~\citep{Bengio+chapter2007}.'').

The corresponding references are to be listed in alphabetical order of
authors, in the \textsc{References} section. As to the format of the
references themselves, any style is acceptable as long as it is used
consistently.

\subsection{Footnotes}

Indicate footnotes with a number\footnote{Sample of the first footnote} in the
text. Place the footnotes at the bottom of the page on which they appear.
Precede the footnote with a horizontal rule of 2~inches
(12~picas).\footnote{Sample of the second footnote}

\subsection{Figures}

All artwork must be neat, clean, and legible. Lines should be dark
enough for purposes of reproduction; art work should not be
hand-drawn. The figure number and caption always appear after the
figure. Place one line space before the figure caption, and one line
space after the figure. The figure caption is lower case (except for
first word and proper nouns); figures are numbered consecutively.

Make sure the figure caption does not get separated from the figure.
Leave sufficient space to avoid splitting the figure and figure caption.

You may use color figures.
However, it is best for the
figure captions and the paper body to make sense if the paper is printed
either in black/white or in color.
\begin{figure}[h]
\begin{center}
\fbox{\rule[-.5cm]{0cm}{4cm} \rule[-.5cm]{4cm}{0cm}}
\end{center}
\caption{Sample figure caption.}
\end{figure}

\subsection{Tables}

All tables must be centered, neat, clean and legible. Do not use hand-drawn
tables. The table number and title always appear before the table. See
Table~\ref{sample-table}.

Place one line space before the table title, one line space after the table
title, and one line space after the table. The table title must be lower case
(except for first word and proper nouns); tables are numbered consecutively.

\begin{table}[t]
\caption{Sample table title}
\label{sample-table}
\begin{center}
\begin{tabular}{ll}
\multicolumn{1}{c}{\bf PART}  &\multicolumn{1}{c}{\bf DESCRIPTION}
\\ \hline \\
Dendrite         &Input terminal \\
Axon             &Output terminal \\
Soma             &Cell body (contains cell nucleus) \\
\end{tabular}
\end{center}
\end{table}

\section{Default Notation}

In an attempt to encourage standardized notation, we have included the
notation file from the textbook, \textit{Deep Learning}
\cite{goodfellow2016deep} available at
\url{https://github.com/goodfeli/dlbook_notation/}.  Use of this style
is not required and can be disabled by commenting out
\texttt{math\_commands.tex}.

\centerline{\bf Numbers and Arrays}
\bgroup
\def1.5{1.5}
\begin{tabular}{p{1in}p{3.25in}}
$\displaystyle a$ & A scalar (integer or real)\\
$\displaystyle {\bm{a}}$ & A vector\\
$\displaystyle {\bm{A}}$ & A matrix\\
$\displaystyle {\tens{A}}$ & A tensor\\
$\displaystyle {\bm{I}}_n$ & Identity matrix with $n$ rows and $n$ columns\\
$\displaystyle {\bm{I}}$ & Identity matrix with dimensionality implied by context\\
$\displaystyle {\bm{e}}^{(i)}$ & Standard basis vector $[0,\dots,0,1,0,\dots,0]$ with a 1 at position $i$\\
$\displaystyle \text{diag}({\bm{a}})$ & A square, diagonal matrix with diagonal entries given by ${\bm{a}}$\\
$\displaystyle {\textnormal{a}}$ & A scalar random variable\\
$\displaystyle {\mathbf{a}}$ & A vector-valued random variable\\
$\displaystyle {\mathbf{A}}$ & A matrix-valued random variable\\
\end{tabular}
\egroup
\vspace{0.25cm}

\centerline{\bf Sets and Graphs}
\bgroup
\def1.5{1.5}

\begin{tabular}{p{1.25in}p{3.25in}}
$\displaystyle {\mathbb{A}}$ & A set\\
$\displaystyle \mathbb{R}$ & The set of real numbers \\
$\displaystyle \{0, 1\}$ & The set containing 0 and 1 \\
$\displaystyle \{0, 1, \dots, n \}$ & The set of all integers between $0$ and $n$\\
$\displaystyle [a, b]$ & The real interval including $a$ and $b$\\
$\displaystyle (a, b]$ & The real interval excluding $a$ but including $b$\\
$\displaystyle {\mathbb{A}} \backslash {\mathbb{B}}$ & Set subtraction, i.e., the set containing the elements of ${\mathbb{A}}$ that are not in ${\mathbb{B}}$\\
$\displaystyle {\mathcal{G}}$ & A graph\\
$\displaystyle \parents_{\mathcal{G}}({\textnormal{x}}_i)$ & The parents of ${\textnormal{x}}_i$ in ${\mathcal{G}}$
\end{tabular}
\vspace{0.25cm}

\centerline{\bf Indexing}
\bgroup
\def1.5{1.5}

\begin{tabular}{p{1.25in}p{3.25in}}
$\displaystyle {a}_i$ & Element $i$ of vector ${\bm{a}}$, with indexing starting at 1 \\
$\displaystyle {a}_{-i}$ & All elements of vector ${\bm{a}}$ except for element $i$ \\
$\displaystyle {A}_{i,j}$ & Element $i, j$ of matrix ${\bm{A}}$ \\
$\displaystyle {\bm{A}}_{i, :}$ & Row $i$ of matrix ${\bm{A}}$ \\
$\displaystyle {\bm{A}}_{:, i}$ & Column $i$ of matrix ${\bm{A}}$ \\
$\displaystyle {\etens{A}}_{i, j, k}$ & Element $(i, j, k)$ of a 3-D tensor ${\tens{A}}$\\
$\displaystyle {\tens{A}}_{:, :, i}$ & 2-D slice of a 3-D tensor\\
$\displaystyle {\textnormal{a}}_i$ & Element $i$ of the random vector ${\mathbf{a}}$ \\
\end{tabular}
\egroup
\vspace{0.25cm}

\centerline{\bf Calculus}
\bgroup
\def1.5{1.5}
\begin{tabular}{p{1.25in}p{3.25in}}
$\displaystyle\frac{d y} {d x}$ & Derivative of $y$ with respect to $x$\\ [2ex]
$\displaystyle \frac{\partial y} {\partial x} $ & Partial derivative of $y$ with respect to $x$ \\
$\displaystyle \nabla_{\bm{x}} y $ & Gradient of $y$ with respect to ${\bm{x}}$ \\
$\displaystyle \nabla_{\bm{X}} y $ & Matrix derivatives of $y$ with respect to ${\bm{X}}$ \\
$\displaystyle \nabla_{\tens{X}} y $ & Tensor containing derivatives of $y$ with respect to ${\tens{X}}$ \\
$\displaystyle \frac{\partial f}{\partial {\bm{x}}} $ & Jacobian matrix ${\bm{J}} \in \mathbb{R}^{m\times n}$ of $f: \mathbb{R}^n \rightarrow \mathbb{R}^m$\\
$\displaystyle \nabla_{\bm{x}}^2 f({\bm{x}})\text{ or }{\bm{H}}( f)({\bm{x}})$ & The Hessian matrix of $f$ at input point ${\bm{x}}$\\
$\displaystyle \int f({\bm{x}}) d{\bm{x}} $ & Definite integral over the entire domain of ${\bm{x}}$ \\
$\displaystyle \int_{\mathbb{S}} f({\bm{x}}) d{\bm{x}}$ & Definite integral with respect to ${\bm{x}}$ over the set ${\mathbb{S}}$ \\
\end{tabular}
\egroup
\vspace{0.25cm}

\centerline{\bf Probability and Information Theory}
\bgroup
\def1.5{1.5}
\begin{tabular}{p{1.25in}p{3.25in}}
$\displaystyle P({\textnormal{a}})$ & A probability distribution over a discrete variable\\
$\displaystyle p({\textnormal{a}})$ & A probability distribution over a continuous variable, or over
a variable whose type has not been specified\\
$\displaystyle {\textnormal{a}} \sim P$ & Random variable ${\textnormal{a}}$ has distribution $P$\\% so thing on left of \sim should always be a random variable, with name beginning with \r
$\displaystyle  \mathbb{E}_{{\textnormal{x}}\sim P} [ f(x) ]\text{ or } \mathbb{E} f(x)$ & Expectation of $f(x)$ with respect to $P({\textnormal{x}})$ \\
$\displaystyle \mathrm{Var}(f(x)) $ &  Variance of $f(x)$ under $P({\textnormal{x}})$ \\
$\displaystyle \mathrm{Cov}(f(x),g(x)) $ & Covariance of $f(x)$ and $g(x)$ under $P({\textnormal{x}})$\\
$\displaystyle H({\textnormal{x}}) $ & Shannon entropy of the random variable ${\textnormal{x}}$\\
$\displaystyle D_{\mathrm{KL}} ( P \Vert Q ) $ & Kullback-Leibler divergence of P and Q \\
$\displaystyle \mathcal{N} ( {\bm{x}} ; {\bm{\mu}} , {\bm{\Sigma}})$ & Gaussian distribution %
over ${\bm{x}}$ with mean ${\bm{\mu}}$ and covariance ${\bm{\Sigma}}$ \\
\end{tabular}
\egroup
\vspace{0.25cm}

\centerline{\bf Functions}
\bgroup
\def1.5{1.5}
\begin{tabular}{p{1.25in}p{3.25in}}
$\displaystyle f: {\mathbb{A}} \rightarrow {\mathbb{B}}$ & The function $f$ with domain ${\mathbb{A}}$ and range ${\mathbb{B}}$\\
$\displaystyle f \circ g $ & Composition of the functions $f$ and $g$ \\
  $\displaystyle f({\bm{x}} ; {\bm{\theta}}) $ & A function of ${\bm{x}}$ parametrized by ${\bm{\theta}}$.
  (Sometimes we write $f({\bm{x}})$ and omit the argument ${\bm{\theta}}$ to lighten notation) \\
$\displaystyle \log x$ & Natural logarithm of $x$ \\
$\displaystyle \sigma(x)$ & Logistic sigmoid, $\displaystyle \frac{1} {1 + \exp(-x)}$ \\
$\displaystyle \zeta(x)$ & Softplus, $\log(1 + \exp(x))$ \\
$\displaystyle || {\bm{x}} ||_p $ & $L^p$ norm of ${\bm{x}}$ \\
$\displaystyle || {\bm{x}} || $ & $L^2$ norm of ${\bm{x}}$ \\
$\displaystyle x^+$ & Positive part of $x$, i.e., $\max(0,x)$\\
$\displaystyle \bm{1}_\mathrm{condition}$ & is 1 if the condition is true, 0 otherwise\\
\end{tabular}
\egroup
\vspace{0.25cm}

\section{Final instructions}
Do not change any aspects of the formatting parameters in the style files.
In particular, do not modify the width or length of the rectangle the text
should fit into, and do not change font sizes (except perhaps in the
\textsc{References} section; see below). Please note that pages should be
numbered.

\section{Preparing PostScript or PDF files}

Please prepare PostScript or PDF files with paper size ``US Letter'', and
not, for example, ``A4''. The -t
letter option on dvips will produce US Letter files.

Consider directly generating PDF files using \verb+pdflatex+
(especially if you are a MiKTeX user).
PDF figures must be substituted for EPS figures, however.

Otherwise, please generate your PostScript and PDF files with the following commands:
\begin{verbatim}
dvips mypaper.dvi -t letter -Ppdf -G0 -o mypaper.ps
ps2pdf mypaper.ps mypaper.pdf
\end{verbatim}

\subsection{Margins in LaTeX}

Most of the margin problems come from figures positioned by hand using
\verb+\special+ or other commands. We suggest using the command
\verb+\includegraphics+
from the graphicx package. Always specify the figure width as a multiple of
the line width as in the example below using .eps graphics
\begin{verbatim}
  \usepackage[dvips]{graphicx} ...
  \includegraphics[width=0.8\linewidth]{myfile.eps}
\end{verbatim}
or
\begin{verbatim}
  \usepackage[pdftex]{graphicx} ...
  \includegraphics[width=0.8\linewidth]{myfile.pdf}
\end{verbatim}
for .pdf graphics.
See section~4.4 in the graphics bundle documentation (\url{http://www.ctan.org/tex-archive/macros/latex/required/graphics/grfguide.ps})

A number of width problems arise when LaTeX cannot properly hyphenate a
line. Please give LaTeX hyphenation hints using the \verb+\-+ command.

\subsubsection*{Author Contributions}
If you'd like to, you may include  a section for author contributions as is done
in many journals. This is optional and at the discretion of the authors.

\subsubsection*{Acknowledgments}
Use unnumbered third level headings for the acknowledgments. All
acknowledgments, including those to funding agencies, go at the end of the paper.
\subsection{Implementation for Atari games}

\textbf{Codebase.}
For Atari games, we use the codebase from \url{https://github.com/Kaixhin/Rainbow} which is an implementation for data-efficient Rainbow \citep{van2019use}.

\textbf{Network architecture.}
To implement our algorithm, we modify the architecture for the value network as shown in Figure \ref{fig:nn_structure} Left.
In data-efficient Rainbow, the state embedding has a dimension of $576$.
We maintain an action embedding for each action, which is a vector of the same dimension and treated as trainable parameters.
Then, we generate the state-action embedding by conducting an element-wise product to the state embedding and the action embedding.
This state-action embedding is shared with the auxiliary task and the main RL task.
Afterwards, the value network outputs the return distribution for this state-action pair (noting that Rainbow uses a distributional RL algorithm, C51 \citep{bellemare2017distributional}) instead of the return distributions of all actions for the input state as is in the original implementation.

\textbf{Hyperparameters.}
We use exactly the same hyperparameters as those used in \citet{van2019use} and CURL \citep{srinivas2020curl} to quantify the gain brought by our auxiliary task and compare with CURL. We refer the readers to their paper for the detailed list of the hyperparameters.

\textbf{Balancing the auxiliary loss and the main RL loss.}
Unlike CURL (or other previous work such as \citet{jaderberg2016reinforcement,yarats2019improving}) that uses different/learned coefficients/learning rates for different games to balance the auxiliary task and the RL updates, our algorithm uses equal weight and learning rate for both the auxiliary task and the main RL task.
This demonstrates the our auxiliary task is robust and does not need careful tuning for these hyperparameters compared with the previous work.

\textbf{Auxiliary loss.}
Since the rewards in Atari games are sparse, we divide the segments such that all the state-action pairs within the same segment have the same return.
This corresponds to the setting of Z-learning with $K\to\infty$ where the positive sample has exactly the same return with that of the anchor.
Then, the auxiliary loss for each update is calculated as follows:
First, we sample a batch of $64$ anchor state-action pairs from the prioritized replay memory.
Then, for each state-action pair, we sample the corresponding positive pair (i.e., the state-action pair within the same segment as the anchor state-action pair) and the corresponding negative pair (randomly selected from the replay memory).
The auxiliary loss is calculated on these samples with effectively $|\mathcal{D}|=128$.

\subsection{Implementation for DMControl suite}

\textbf{Codebase.}
We use SAC as the base RL algorithm and build our algorithm on top of the publicly released implementation from CURL \citep{srinivas2020curl}.

\textbf{Network architecture.}
Similarly, we modify the architecture for the critic network in SAC.
In SAC, the state embedding has a dimension of 50.
Since the actions are continuous vectors of dimension $d$ in the continuous control tasks of DMControl suite, we directly concatenate the action to the state embedding, resulting in a state-action embedding of size $50+d$.
Then, the critic network receives the state-action embedding as the input and outputs the Q value.
The actor network receives the state embedding as the input and output the action selection distribution on the corresponding state.
Note that, although our auxiliary loss is based on the state-action embedding, the state embedding used by the actor network is also trained by the auxiliary loss through back-propagation of the gradients.

\textbf{Hyperparameters.}
We set the threshold for dividing the segments to $1.0$, i.e., when appending transitions to the replay buffer, we start a new segment when the cumulative reward within the last segment exceeds this threshold.
The auxiliary loss and the hyperparameters to balance the auxiliary loss and the main RL loss are the same as those used for Atari games.
Other hyperparameters we use are exactly the same as those in CURL implementation and we refer the readers to their paper for the details.

\section{Introduction}

\iffalse
End-to-end model-free reinforcement learning (RL) directly optimizes different components simultaneously for the end objective and has achieved many successes in complex control tasks recently \citep{mnih2015human,haarnoja2018soft,hwangbo2019learning}.
However, deep RL agents still need massive number of interactions with the environment to learn a good policy.
One reason for the sample inefficiency of RL algorithms is that the policy and the intermediate representation are trained simultaneously based on weak signals provided by the end objective, such as the gradient of cumulative reward \citep{schulman2015trust,schulman2017proximal}, or the TD error \citep{mnih2015human,hessel2017rainbow}.
The signals become even weaker when the reward is sparse, delayed or noisy.
Since a good representation is a prerequisite for a good policy, the weak signals result in slow learning for the representation and therefore the policy.

The \emph{auxiliary task} dating back to \citet{suddarth1990rule} can be used to improve the sample efficiency.
Auxiliary tasks are additional tasks that are learned simultaneously with the main RL agent to improve the quality of the shared representation.
The improved representation can help the agent to better generalize from one state to another and therefore improve the sample efficiency.

One important ingredient for designing an auxiliary task is the source of information used.
Some auxiliary tasks exploit the prior knowledge associated with the state representation or the dynamics.
For examples, auxiliary tasks can be designed to exploit the invariance properties of the images-based states (e.g., the occlusion invariance \citep{purushwalkam2020demystifying}) such as
CURL \citep{srinivas2020curl} and CPC \citep{oord2018representation}.
Also, auxiliary tasks can be designed to exploit the prior knowledge of the dynamics of the main task \citep[e.g., ][]{shen2020auxiliary}.
Although such auxiliary tasks achieve good performance in certain types of RL tasks, they are limited by the prior that they rely on and therefore may not applicable to other scenarios.
Other auxiliary tasks solely based on the structure of Markov decision process (MDP) such as learning a bisimulation model \citep{gelada2019deepmdp,zhang2020learning}, approximating other value functions \citep{bellemare2019geometric,dabney2020value} or building a world model \citep{franccois2019combined,hafner2019dream,hafner2019learning}.
In this paper, we propose to use a new form of information based on the structure of MDP, the distribution over the returns, to design the auxiliary task.
Such information is available in different RL tasks and easy to extract (using only the samples from the standard RL algorithms).

Another essential ingredient for designing an auxiliary task is the way we integrate the information to the representation.
Standard RL algorithms learn the policy conservatively \citep{kakade2002approximately} or update the value function in a bootstrapped way \citep{mnih2015human}, resulting in slow learning.
Since a good policy is built upon a good representation, we hope to learn the representation via the auxiliary task in a fast way
(e.g., with a relatively fixed target from the start).
From this perspective, contrastive learning is good candidate where
the representation is learned by minimizing the embedding distance between similar instances (e.g., states or state-actions).
Compared with the values or the policies for each instance, the similarity between the instances is relative fixed which enables a fast learning for the representation.
One successful example of contrastive learning based auxiliary task is CURL \citep{srinivas2020curl} where a different source of information is used compared with ours (cf. Figure \ref{fig:fig1}).

\begin{wrapfigure}{r}{0.6\textwidth}
\vspace{-20pt}
  \begin{center}
    \includegraphics[width=0.6\textwidth]{figures/fig1.png}
  \end{center}
  \label{fig:fig1}
  \caption{
  We use a different source of information for contrastive learning based auxiliary task from the previous work.
  CURL minimizes the embedding distance between different augmentations of the same state (top), whereas our auxiliary task aggregates the states with similar distributions over returns (bottom).
  }
\end{wrapfigure}

In this paper, we deduce our contrastive learning based auxiliary task from a perspective of the state-action abstraction.
We define a new form of state-action abstraction, $Z^\pi$ irrelevance, that aggregates the state-actions with the same distribution over the returns.
We design an algorithm, Z-learning, which is mainly based on a contrastive loss, and prove that the algorithm can learn $Z^\pi$ irrelevance provably efficiently.
Then, we adapt Z-learning as an auxiliary task and proposed Return based Contrastive learning for RL (RCRL).
Our auxiliary task is well-aligned with the standard RL algorithm and boosts the training process by learning a representation that can better generalize across the states.
We conduct experiments on DeepMind Control Suite \citep{tassa2018deepmind} and Atari games \citep{bellemare2013arcade}.
The experiment results show that our auxiliary task combined with Rainbow \citep{hessel2017rainbow} for discrete control tasks or SAC \citep{haarnoja2018soft} for continuous control tasks achieves superior performance over other baseline auxiliary tasks.

Our contributions are summarized as follows:
\begin{itemize}
    \item (Section \ref{sec:preliminary}) We introduce a new form of state-action abstraction, $Z^\pi$ irrelevance and identify its properties.
    \item (Section \ref{sec:method}) We propose Z-learning to learn $Z^\pi$ irrelevance probably efficiently.
    \item (Section \ref{sec:method}) We propose SCRL that adapts Z-learning as an auxiliary task to boost the learning for standard RL algorithms.
    \item (Section \ref{sec:experiment}) We conduct experiments on Procgen, DMC and Atari, and demonstrate the superior performance of SCRL.
\end{itemize}

\fi

Deep reinforcement learning~(RL) algorithms can learn representations from high-dimensional inputs,
{as well as learn policies} based on such representations to maximize long-term returns simultaneously.
However,
{deep RL algorithms} typically require
large numbers of samples, which can be quite expensive to obtain
~\citep{mnih2015human}.
In contrast, it is usually much more sample efficient to learn
{policies} with learned representations/extracted features~\citep{srinivas2020curl}.
To this end, various auxiliary tasks have been proposed to accelerate representation learning
{in aid of}
the main RL task ~\citep{suddarth1990rule,sutton2011horde,gelada2019deepmdp,bellemare2019geometric,franccois2019combined,shen2020auxiliary,zhang2020learning,dabney2020value,srinivas2020curl}.
Representative examples of auxiliary tasks include predicting the future in either the pixel space or the latent space with reconstruction-based losses~\citep[e.g., ][]{jaderberg2016reinforcement,hafner2019dream,hafner2019learning}.

Recently, contrastive learning has been introduced to construct auxiliary tasks and achieves better performance compared to reconstruction based methods in accelerating RL algorithms~\citep{oord2018representation,srinivas2020curl}.
Without the need to reconstruct inputs such as raw pixels, contrastive learning based methods can ignore irrelevant features such as static background in games and learn more compact representations.
\cite{oord2018representation} propose a contrastive representation learning method based on the temporal structure of state sequence.
\cite{srinivas2020curl} propose to leverage
the prior knowledge from computer vision, learning representations that are
{invariant to image augmentation.}
{However, existing works mainly construct contrastive auxiliary losses
in an unsupervised manner,
without considering feedback signals in RL problems as supervision.
}

In this paper, we take a further step to
{leverage the return feedback}
to design a contrastive auxiliary loss to accelerate RL algorithms.
Specifically, we propose a novel method, called Return-based Contrastive representation learning for Reinforcement Learning (RCRL).
In our method, given an anchor state-action pair, we choose
a state-action pair with the same or similar return as the positive sample, and a state-action pair with different return as the negative sample.
{Then, we train a discriminator to classify between positive and negative samples given the anchor based on their representations as the auxiliary task.}
The intuition here is to learn state-action representations that capture return-relevant features while \textcolor{black}{ignoring} return-irrelevant features.

{From a theoretical perspective, RCRL is supported by a novel state-action abstraction, called $Z^\pi$-irrelevance.}
$Z^\pi$-irrelevance abstraction aggregates
state-action pairs
{with similar return distributions}
under certain policy $\pi$.
{We show that $Z^\pi$-irrelevance abstraction can reduce the size of the state-action space (cf. Appendix \ref{app:irrelevance}) as well as
approximate the Q values arbitrarily accurately
(cf. Section \ref{sec:Zpi-irrelevance}).}
We further propose a method called Z-learning that can calculate $Z^\pi$-irrelevance abstraction with sampled returns rather than the return distribution, which is hardly available in practice.
{Z-learning can learn $Z^\pi$-irrelevance abstraction provably efficiently.}
Our algorithm {RCRL can be seen as the empirical version of Z-learning by making a few approximations such as
integrating with deep RL algorithms, and collecting positive pairs
within a consecutive segment in a trajectory of the anchors. }

We conduct experiments on Atari games \citep{bellemare2013arcade} and DeepMind Control suite \citep{tassa2018deepmind} in low data regime.
The experiment results show that our auxiliary task combined with Rainbow \citep{hessel2017rainbow} for discrete control tasks or SAC \citep{haarnoja2018soft} for continuous control tasks achieves superior performance over other
state-of-the-art baselines for this regime.
Our method can be further combined with existing unsupervised contrastive learning methods to achieve even better performance.
We also perform a detailed analysis on how the representation changes during training with/without our auxiliary loss.
We find that a good embedding network assigns similar/dissimilar representations to state-action pairs with similar/dissimilar return distributions, and our algorithm can boost such generalization and speed up training.

Our contributions are summarized as follows:
\begin{itemize}
    \item We introduce a novel contrastive loss based on return, to learn return-relevant representations and speed up deep RL algorithms.

    \item We theoretically build the connection between the contrastive loss and a new form of state-action abstraction, which can reduce the size of the state-action space as well as approximate the Q values arbitrarily accurately.

    \item Our algorithm achieves superior performance against strong baselines in Atari games and DeepMind Control suite \textcolor{black}{in low data regime}. Besides, the performance can be further enhanced when combined with existing auxiliary tasks
\end{itemize}
\subsection{Comparison with bisimulation.}

We consider a bisimulation $\phi$ and denote the set of abstract states as $\mathcal{Z}:=\phi(\mathcal{S})$.
The bisimulation $\phi_B$ is defined as follows:

\begin{definition}[Bisimulation \citep{givan2003equivalence}]
$\phi_B$ is bisimulation if
$\forall s_1, s_2 \in \mathcal{S}$ where $\phi_B(s_1)=\phi_B(s_2)$, $\forall a \in \mathcal{A}, z'\in \mathcal{Z}$,
\begin{equation*}
    R(s_1, a) = R(s_2, a), \quad
    \sum_{s'\in\phi_B^{-1}(z')} P(s'|s_1, a) = \sum_{s'\in\phi_B^{-1}(z')}
    P(s'|s_2, a).
\end{equation*}
\end{definition}

Notice that the coarseness of $Z^\pi$-irrelevance is dependent on $K$, the number of bins for the return.
When $K\to\infty$, two state-action pairs $x$ and $x'$ are aggregated only when $Z^\pi(x) \stackrel{D}{=} Z^\pi(x')$, which is a strict condition resulting in a fine-grained abstraction.
Here, we provide the following proposition to illustrate that $Z^\pi$-irrelevance abstraction is coarser than bisimulation even when $K\to\infty$.

\begin{proposition}
\label{prop:formal_coarser}
Given ${\phi}_B$ to be the coarsest bisimulation, ${\phi}_B$ induces a $Z^\pi$-irrelevance abstraction for any policy $\pi$ defined over $\mathcal{Z}$. Specifically,
if
$\forall s_1, s_2 \in \mathcal{S}$ where
${\phi}_B(s_1)=\phi_B(s_2)$, then $Z^\pi(s_1,a) \stackrel{D}{=} Z^\pi(s_2, a), \forall a\in\mathcal{A}$ for any policy $\pi$ defined over $\mathcal{Z}$.
\end{proposition}

Consider a state-action abstraction $\tilde{\phi}_B$ that is augmented from the coarsest bisimulation $\phi_B$:
$\tilde{\phi}_B(s_1,a_1) = \tilde{\phi}_B(s_2,a_2)$ if and only if $\phi_B(s_1) = \phi(s_2)$ and $a_1 = a_2$.
The proposition indicates that $|\phi_B(\mathcal{S})|\,|\mathcal{A}|=|\tilde{\phi}_B(\mathcal{X})| \ge N_{\pi,\infty} \ge N_{\pi,K}$ for any $K$ and for any $\pi$ defined over $\mathcal{Z}$, i.e., bisimulation is no coarser than $Z^\pi$-irrelevance.
Note that there exists an optimal policy that is defined over $\mathcal{Z}$ \citep{li2006towards}.
In practice, $Z^\pi$-irrelevance should be far more coarser than bisimulation when we only consider one specific policy $\pi$.
Therefore learning a $Z^\pi$-irrelevance should be easier than learning a bisimulation.

\begin{proof}
First, for a fixed policy $\pi: \mathcal{Z} \to \Delta^\mathcal{A}$, we prove that
if two state distributions projected to $\mathcal{Z}$ are identical,
the corresponding reward distributions or the next state distributions projected to $\mathcal{Z}$
will be identical.
Then, we use such invariance property to
prove the proposition.

The proof for the invariance property goes as follows:
Consider two identical state distributions over $\mathcal{Z}$ on the $t$-th step, $P_1, P_2 \in \Delta^{\mathcal{Z}}$ such that $P_1 = P_2$.
Notice that we only require that the two state distributions are identical when projected to $\mathcal{Z}$, so therefore they may be different on $\mathcal{S}$.
Specifically, if we denote the state distribution over $\mathcal{S}$ as $P(s) = P(z)q(s|z)$ where $\phi_B(s) = z$,
the distribution $q$ for the two state distributions may be different.
However, we will show that this is sufficient to ensure that the corresponding state distributions on $\mathcal{Z}$ are identical on the next step (and therefore the subsequent steps).

We denote $R_1$ and $R_2$ as the reward on the $t$-th step, which are random variables.
The reward distribution on the $t$-th step is specified as follows:
\begin{equation}
\begin{aligned}
P(R_1 = r | P_1) = & \sum_{z\in\mathcal{Z},a\in\mathcal{A},s\in\mathcal{S}} P(r|s,a) q_1(s|z) \pi(a|z) P_1(z) \\
= & \sum_{z\in\mathcal{Z},a\in\mathcal{A}} P(r|z,a) \pi(a|z) P_1(z) \\
= & \sum_{z\in\mathcal{Z},a\in\mathcal{A}} P(r|z,a) \pi(a|z) P_2(z) \\
= & P(R_2 = r | P_2),
\end{aligned}
\end{equation}
for any $r\in\mathbb{R}$, where we use the property of bisimulation, $R|s,a \stackrel{D}{=} R|s',a, \forall \phi_B(s) = \phi_B(s')$, in the second equation. This indicates that
$R_1|P_1 \stackrel{D}{=} R_2|P_2$.

Similarly, we denote $P'_1$ and $P'_2$ as the state distribution over $\mathcal{Z}$ on the $(t+1)$-th step. We have
\begin{equation}
\begin{aligned}
P'_1(z')| P_1 = & \sum_{z\in\mathcal{Z},a\in\mathcal{A}}
\sum_{s\in\mathcal{S}}
\left( \sum_{s'\in\phi_B^{-1}(z')} P(s'|s,a) \right)
q_1(s|z) \pi(a|z) P_1(z) \\
= & \sum_{z\in\mathcal{Z},a\in\mathcal{A}} P(z'|z,a) \pi(a|z) P_1(z) \\
= & \sum_{z\in\mathcal{Z},a\in\mathcal{A}} P(z'|z,a) \pi(a|z) P_2(z) \\
= & P'_2(z')| P_1,
\end{aligned}
\end{equation}
for any $z'\in\mathcal{Z}$, where we also use the property of bisimulation in the second equation, i.e., the quantity in the bracket equals for different $s$ such that $q(s|z) > 0$ with a fixed $z$. This indicates that $P'_1|P_1 = P'_2|P_2$.

At last,
consider two state $s_1, s_2\in\mathcal{S}$ with $\phi_B(s_1)=\phi_B(s_2)=z$ on the first step, and we take $a\in\mathcal{A}$ for both states on this step and later take the actions following $\pi$.
We denote the subsequent reward and state distribution over $\mathcal{Z}$ as $R_1^{(1)}, P_1^{(2)}, R_1^{(2)}, P_1^{(3)}, \cdots $ and $R_2^{(1)}, P_2^{(2)}, R_2^{(2)}, P_2^{(3)}, \cdots $ respectively, where the superscripts indicate time steps.
Therefore, we can deduce that $P_1^{(1)} = P_2^{(1)}, R_1^{(1)} \stackrel{D}{=} R_2^{(1)}, P_1^{(2)} = P_2^{(2)}, R_1^{(2)} \stackrel{D}{=} R_2^{(2)}, \cdots $ and therefore $Z^\pi(s_1, a) \stackrel{D}{=} Z^\pi(s_2, a), \forall a \in \mathcal{A}$.

\end{proof}

\subsection{\textcolor{black}{Comparison with $\pi$-bisimulation}}

\textcolor{black}{Similar to $Z^\pi$-irrelevance, a recently proposed state abstraction $\pi$-bisimulation \citep{castro2020scalable} is also tied to a behavioral policy $\pi$.
It is beneficial to compare the coarseness of $Z^\pi$-irrelevance and $\pi$-bisimulation.
For completeness, we restate the definition of $\pi$-bisimulation.}

\textcolor{black}{
\begin{definition}[$\pi$-bisimulation \citep{castro2020scalable}]
Given a policy $\pi$, $\phi_{B, \pi}$ is a $\pi$-bisimulation if $\forall s_1, s_2 \in \mathcal{S}$ where $\phi_{B, \pi}(s_1) = \phi_{B, \pi}(s_2)$,
\begin{equation*}
\begin{aligned}
\sum_a \pi(a|s_1) R(s_1, a) & = \sum_a \pi(a|s_2) R(s_2, a) \\
\sum_a \pi(a|s_1) \sum_{s'\in \phi_{B, \pi}^{-1}(z')} P(s'|s_1,a) & = \sum_a \pi(a|s_2) \sum_{s'\in \phi_{B, \pi}^{-1}(z')} P(s'|s_2,a), \qquad \forall z' \in \mathcal{Z},
\end{aligned}
\end{equation*}
where $\mathcal{Z}:=\phi_{B, \pi}(\mathcal{S})$ is the set of abstract states.
\end{definition}
}

\textcolor{black}{However, $\pi$-bisimulation does not consider the state-action pair that is not visited under the policy $\pi$ (e.g., a state-action pair $(s,a)$ when $\pi(a|s)=0$), whereas $Z^\pi$-irrelevance is defined on all the state-action pairs.
Therefore, it is hard to build connection between them unless we also define $Z^\pi$-irrelevance on the state space (instead of the state-action space) in a similar way.}

\textcolor{black}{
\begin{definition}[$Z^\pi$-irrelevance on the state space]
Given $s\in\mathcal{S}$, we denote $\mathbb{Z}^\pi(s) := \sum_a \pi(a|s) \mathbb{Z}^\pi(s,a)$.
Given a policy $\pi$, $\phi$ is a $Z^\pi$-irrelevance if $\forall s_1, s_2 \in \mathcal{S}$ where $\phi(s_1) = \phi(s_2)$,
$\mathbb{Z}^\pi(s_1) = \mathbb{Z}^\pi(s_2)$.
\end{definition}
}

\textcolor{black}{Based on the above definitions, the following proposition indicates that such $Z^\pi$-irrelevance is coarser than $\pi$-bisimulation.}

\textcolor{black}{
\begin{proposition}
Given a policy $\pi$ and $\phi_{B, \pi}$ to be the coarsest $\pi$-bisimulation, if $\forall s_1, s_2\in \mathcal{S}$ where $\phi_{B, \pi}(s_1)=\phi_{B, \pi}(s_2)$, then $\mathbb{Z}^\pi(s_1) = \mathbb{Z}^\pi(s_2)$.
\end{proposition}
}
\begin{proof}
\textcolor{black}{
Starting from $s_1$ and $s_2$ and following the policy $\pi$, the reward distribution and the state distribution over $\mathcal{Z}$ on each step are identical, which can be proved by induction. Then, we can conclude that $Z^\pi(s_1) \stackrel{D}{=} Z^\pi(s_2)$ and thus $\mathbb{Z}^\pi(s_1) = \mathbb{Z}^\pi(s_2)$.
}
\end{proof}

\subsection{Bound on the representation error}

For completeness, we restate Proposition \ref{prop:aligned}.

\begin{proposition}
Given a policy $\pi$ and any $Z^\pi$-irrelevance $\phi:\mathcal{X}\to [N]$, there exists a function ${Q}: [N] \to \mathbb{R}$ such that $|{Q}({\phi}(x)) - Q^\pi(x)| \le \frac{R_{\max}-R_{\min}}{K}, \forall x \in \mathcal{X}$.
\end{proposition}

\begin{proof}
Given a policy $\pi$ and a $Z^\pi$ irrelevance $\phi$, we can construct a ${Q}$ such that ${Q}({\phi}(x)) = Q^\pi(x), \forall x \in \mathcal{X}$ in the following way:
For all $x\in\mathcal{X}$, one by one, we assign $Q(z) \leftarrow Q^\pi(x)$, where $z = \phi(x)$.
In this way, for any $x\in\mathcal{X}$ with $z=\phi(x)$, $Q(z)=Q^\pi(x')$ for some $x'\in\mathcal{X}$ such that $\mathbb{Z}^\pi(x')=\mathbb{Z}^\pi(x)$.
This implies that $|Q(z)-Q^\pi(x)|=|Q^\pi(x')-Q^\pi(x)|\le\frac{R_{\max} - R_{\min}}{K}$.
This also applies to the optimal policy $\pi^*$.
\end{proof}
\section{Related Work}

\input{related_work_new}

\section{Preliminary}
\label{sec:preliminary}

\input{preliminary}

\section{Methodology}
\label{sec:method}

\input{method}

\section{Experiment}
\label{sec:experiment}

\input{experiment_zch}

\vspace{-10pt}
\section{Conclusion}
\input{conclusion}

\section{Acknowledgment}
\input{acknowledgments}

\subsection{$Z^\pi$-irrelevance Abstraction}
\label{sec:Zpi-irrelevance}
A state-action abstraction aggregates the state-action pairs with similar properties, resulting in an abstract state-action space denoted as $[N]$,
{where $N$ is the size of abstract state-action space}.
In this paper, we consider a new form of abstraction, $Z^\pi$-irrelevance, defined as follows:
Given a policy $\pi$, $Z^\pi$-irrelevance abstraction is denoted as $\phi: \mathcal{X}\to[N]$ such that,
for any $x_1, x_2 \in \mathcal{X}$ with $\phi(x_1)=\phi(x_2)$, we have $\mathbb{Z}^\pi(x_1) = \mathbb{Z}^\pi(x_2)$.
Given a policy $\pi$ and the parameter for return discretization $K$, we use
$N_{\pi,K}$ to denote the minimum $N$ such that a $Z^\pi$-irrelevance exists.
It is true that
$N_{\pi,K}\le N_{\pi,\infty} \le |\phi_B(\mathcal{S})|\, |\mathcal{A}|$ for any $\pi$ and $K$,
where $|\phi_B(\mathcal{S})|$ is the number of abstract states for the coarsest bisimulation (cf. Appendix \ref{app:irrelevance}).

\begin{proposition}
\label{prop:aligned}
Given a policy $\pi$ and any $Z^\pi$-irrelevance $\phi:\mathcal{X}\to [N]$, there exists a function ${Q}: [N] \to \mathbb{R}$ such that $|{Q}({\phi}(x)) - Q^\pi(x)| \le \frac{R_{\max}-R_{\min}}{K}, \forall x \in \mathcal{X}$.
\end{proposition}
We provide a proof in Appendix \ref{app:irrelevance}.
Note that $K$ controls the coarseness of the abstraction.
When $K\to\infty$, $Z^\pi$-irrelevance can accurately represent the value function and therefore the optimal policy when $\pi\to\pi^*$.
When using an auxiliary task to learn such abstraction, this proposition indicates
that the auxiliary task (to learn a $Z^\pi$-irrelevance) is well-aligned with the main RL task (to approximate $Q^*$).
However, large $K$ results in a fine-grained abstraction which requires us to use a large $N$ and more samples to learn the abstraction (cf. Theorem \ref{theorem:main}).
In practice, this may not be a problem since we learn a state-action representation in a low-dimensional space $\mathbb{R}^d$ instead of $[N]$ and reuse the samples collected by the base RL algorithm.
Also, we do not need to choose a $K$ explicitly in the practical algorithm (cf. Section \ref{sec:RCRL}).

\iffalse
We provide proof in Appendix \ref{app:irrelevance}.
The proposition indicates that, when $K\to\infty$, any $Z^{\pi}$-irrelevance $\phi$ can represent the exact value function for the current policy.
Moreover, considering $\pi=\pi^*$, $\phi$ can represent $Q^*$ and therefore the optimal policy.
Therefore, $\phi$ may serve as a good representation for value functions.
In practice, we learn a state-action representation with the help of an auxiliary task to embed state-action pairs into a low-dimensional space $\mathbb{R}^d$ instead of $[N]$.
We hope that the representation captures the structure of $Z^\pi$-irrelevance for the current policy $\pi$ during the learning, i.e., state-action pairs with similar distributions over returns share similar representations.
In this case, the proposition demonstrates the consistency between the auxiliary task (to learn a $Z^\pi$-irrelevance) and the main RL task (to approximate $Q^*$).

\textcolor{red}{Change the claim for keep optimal policy invariant.}
\fi

\subsection{Z-Learning}

\begin{algorithm}[t]
\caption{Z-learning}
\label{algo:z-learning}
\begin{algorithmic}[1]
\STATE Given the policy $\pi$, the number of bins for the return $K$, a constant $N \ge N_{\pi, K}$,
the encoder class $\Phi_N$, the regressor class $\mathcal{W}_N$,
and a distribution $d\in\Delta^\mathcal{X}$ with $\text{supp}(d)=\mathcal{X}$
\STATE $\mathcal{D}=\emptyset$
\FOR {$i=1, \cdots, n$}
    \STATE $x_1, x_2 \sim d$
    \STATE $R_1 \sim Z^\pi(x_1), R_2 \sim Z^\pi(x_2)$
    \STATE $\mathcal{D} = \mathcal{D} \cup \{(x_1, x_2, y=\mathbb{I}[b(R_1) \neq b(R_2)])\}$
\ENDFOR
\STATE \textcolor{black}{$(\hat{\phi}, \hat{w}) = \text{arg} \min_{\phi \in \Phi_N, w \in \mathcal{W}_N} \mathcal{L}(\phi, w; \mathcal{D})$, where $\mathcal{L}(\phi, w;\mathcal{D})$ is defined in (\ref{eq:loss_function})}
\STATE \textbf{return } the encoder $\hat{\phi}$
\end{algorithmic}
\end{algorithm}

We propose Z-learning to learn
$Z^\pi$-irrelevance based
on a dataset $\mathcal{D}$ with a contrastive loss (see Algorithm \ref{algo:z-learning}).
Each tuple in the dataset is collected as follows:
First, two state-action pairs are drawn i.i.d. from a distribution $d\in\Delta^\mathcal{X}$ with $\text{supp}(d)=\mathcal{X}$ (cf. Line 4 in Algorithm \ref{algo:z-learning}).
In practice, we can sample state-action pairs from the rollouts generated by the policy $\pi$.
In this case, a stochastic policy (e.g., using $\epsilon$-greedy) with a standard ergodic assumption on MDP ensures $\text{supp}(d)=\mathcal{X}$.
Then, we obtain returns for the two state-action pairs (i.e., the discounted sum of the rewards after $x_1$ and $x_2$) which can be obtained by rolling out using the policy $\pi$ (cf. Line 5 in Algorithm \ref{algo:z-learning}).
The binary label $y$ for this state-action pair indicates whether the two returns belong to the same bin (cf. Line 6 in Algorithm \ref{algo:z-learning}).
The contrastive loss is defined as follows:
\begin{equation}
\label{eq:loss_function}
    \min_{
    \phi \in \Phi_N, w\in\mathcal{W}_N
    } \mathcal{L}(\phi, w;\mathcal{D}) := \mathbb{E}_{(x_1, x_2, y)\sim \mathcal{D}} \left[ (w(\phi(x_1), \phi(x_2)) - y)^2 \right],
\end{equation}
\textcolor{black}{
where the class of encoders that map the state-action pairs to $N$ discrete abstractions is defined as $\Phi_N:=\{\mathcal{X}\to[N]\}$,
and the class of tabular regressors is defined as $\mathcal{W}_N:=\{[N]\times[N]\to[0,1]\}$.}
Notice that we choose $N\ge N_{\pi,K}$ to ensure that a $Z^\pi$-irrelevance $\phi:\mathcal{X}\to [N]$ exists.
Also, to aggregate the state-action pairs, $N$ should be smaller than $|\mathcal{X}|$ (otherwise we will obtain an identity mapping).
In this case, mapping two state-action pairs with different return distributions to the same abstraction will increase the loss and therefore is avoided.
The following theorem shows that Z-learning can learn $Z^\pi$-irrelevance provably efficiently.

\begin{theorem}
\label{theorem:main}
Given the encoder $\hat{\phi}$ returned by Algorithm \ref{algo:z-learning}, the following inequality holds with probability $1-\delta$
\textcolor{black}{and for any $x'\in\mathcal{X}$}:
\begin{equation}
\begin{aligned}
\mathbb{E}_{x_1 \sim d, x_2\sim d} & \Big[ \mathbb{I}[\hat{\phi}(x_1)=\hat{\phi}(x_2)] \, \Big|
 \mathbb{Z}^\pi(x')^T \left( \mathbb{Z}^\pi(x_1) - \mathbb{Z}^\pi(x_2) \right)
\Big| \Big] \\
& \quad \quad \le \sqrt{\frac{8N}{n} \Big( 3 + 4N^2 \ln n + 4 \ln |\Phi_N| + 4 \ln (\frac{2}{\delta}) \Big) },
\end{aligned}
\end{equation}
where $|\Phi_N|$ is the cardinality of encoder function class and $n$ is the size of the dataset.
\end{theorem}

We provide the proof in Appendix \ref{app:proof}.
\textcolor{black}{
Although $|\Phi_N|$ is upper bounded by $N^{|\mathcal{X}|}$, it is generally much smaller for deep encoders that generalize over the state-action space.
}
The theorem shows that whenever $\hat{\phi}$ maps two state-actions $x_1, x_2$ to the same abstraction, $\mathbb{Z}^\pi(x_1) \approx \mathbb{Z}^\pi(x_2)$ up to an error proportional to $1/\sqrt{n}$ (ignoring the logarithm factor).
The following corollary shows that $\hat{\phi}$ becomes a $Z^\pi$-irrelevance when $n\to\infty$.

\begin{corollary}
\label{corollary:main}
The encoder $\hat{\phi}$ returned by Algorithm \ref{algo:z-learning} with $n\to\infty$ is a $Z^\pi$-irrelevance,
i.e.,
for any $x_1, x_2\in\mathcal{X}$, $\mathbb{Z}^\pi(x_1)
= \mathbb{Z}^\pi(x_2)$ if $\hat{\phi}(x_1)=\hat{\phi}(x_2)$.
\end{corollary}

\subsection{Return-based contrastive learning for RL (RCRL)}
\label{sec:RCRL}

We adapt Z-learning as an auxiliary task that helps the agent to learn a representation with meaningful semantics.
The auxiliary task based RL algorithm is called RCRL and shown in Algorithm \ref{algo:auxiliary_task}.
Here, we use Rainbow (for discrete control) and SAC (for continuous control) as the base RL algorithm for RCRL.
However, RCRL can also be easily incorporated with other model-free RL algorithms.
While Z-learning relies on a dataset sampled by rolling out the current policy,
RCRL constructs such a dataset using the samples collected by the base RL algorithm and therefore does not require additional samples, e.g., directly using the replay buffer in Rainbow or
SAC
(see Line 7 and 8 in Algorithm \ref{algo:auxiliary_task}).
Compared with Z-learning, we use the state-action embedding network that is shared with the base RL algorithm $\phi:\mathcal{X}\to\mathbb{R}^d$
as the encoder,
and use an additional discriminator trained by the auxiliary task $w:\mathbb{R}^d\times \mathbb{R}^d \to [0,1]$ as the regressor.

However, when implementing Z-learning as the auxiliary task, the labels in the dataset may be unbalanced.
Although this does not cause problems in the theoretical analysis since we assume the Bayes optimizer can be obtained for the contrastive loss,
{it may prevent the discriminator from learning properly in practice (cf. Line 8 in Algorithm \ref{algo:z-learning}).}
To solve this problem, instead of drawing samples independently from the replay buffer $\mathcal{B}$~(analogous to sampling from the distribution $d$ in Z-learning),
we sample the pairs for $\mathcal{D}$ as follows:
As a preparation, we cut the trajectories in $\mathcal{B}$ into segments, where each segment contains state-action pairs with the same or similar returns.
\textcolor{black}{
Specifically, in Atari games, we create a new segment once the agent receives a non-zero reward.
In DMControl tasks, we first prescribe a threshold and then create a new segment once the cumulative reward within the current segment exceeds this threshold.
}
For each sample in $\mathcal{D}$, we first draw an anchor state-action pair from $\mathcal{B}$ randomly.
Afterwards,
we generate a positive sample by
drawing a state-action pair from the same segment of the anchor state-action pair.
Then, we draw
another state-action pair randomly
from $\mathcal{B}$ and use it as the negative sample.

\begin{algorithm}[thbp]
\label{algo:auxiliary_task}
\SetAlgoLined
    Initialize the embedding $\phi_\theta:\mathcal{X} \to \mathbb{R}^d$ and a discriminator $w_\vartheta: \mathbb{R}^d \times \mathbb{R}^d \to [0,1]$

    Initialize the parameters $\varphi$ for the base RL algorithm

    that uses the learned embedding $\phi_\theta$

    Given a batch of samples $\mathcal{D}$, the loss function for the base RL algorithm is $\mathcal{L}_\text{RL}(\phi_\theta, \varphi; \mathcal{D})$

    A replay buffer $\mathcal{B} = \emptyset$

    \ForEach{iteration}{
        Rollout the current policy and store the samples to the replay buffer $\mathcal{B}$

        \textcolor{black}{Draw a batch of samples $\mathcal{D}$ from the replay buffer $\mathcal{B}$}

        \textcolor{black}{Update the parameters with the loss function
        $\mathcal{L}(\phi_\theta, w_\vartheta; \mathcal{D}) + \mathcal{L}_\text{RL}(\phi_\theta, \varphi; \mathcal{D})$}

    }
    \Return{The learned policy}
\caption{Return based Contrastive learning for RL (RCRL)}
\end{algorithm}

\iffalse

\begin{figure}[htbp]
    \centering
    \includegraphics[width=0.6\textwidth]{figures/figure3.png}
    \caption{
  We cut the trajectories into segments where the state-actions in each segment share the same returns.
  The two states are semantically different (e.g., the optimal actions on these two states are opposite) but quite similar in pixels.
  By discriminating the corresponding state-actions with different distributions over returns,
  we attain representation with better generalization and therefore boost the learning.
    }
    \label{fig:fig3}
\end{figure}

We believe the auxiliary task may boost the learning due to the following reasons:
First, by aggregating the state-action space in a soft way, we learn a useful generalization over the state-actions which improves the sample efficiency.
For example, in Montezuma's Revenge, when the agent is standing on a ladder, the only effective actions are climbing up or climbing down, and the remaining actions take no effect.
By aggregating the states with ineffective actions (such that the their representations are similar), we can avoid too much useless exploration and therefore improve sample efficiency.
Another example is shown in Figure \ref{fig:fig3}.
Climbing up is a better action before the agent obtains the key with a reward (State 1), whereas climbing down is a better action after the agent obtains the key (State 2).
However, these two states are quite similar in pixels.
By discriminating these two states, we avoid over generalization between State 1 and State 2.
Second, instead of learning the representation using the TD error for the Q values, such an auxiliary task incorporate more information (e.g., the distribution of the returns and pairwise comparison between the state-actions) into the representation.
Third, the update targets are fast changing due to bootstrapping.
This may lead to a slow learning process for the representation.
The auxiliary task aggregates the state with similar distribution over returns, resulting in a relatively fixed target.
\textcolor{red}{I will revise this paragraph later.}

\fi
{
We believe our auxiliary task may boost the learning, due to better return-induced representations that facilitate generalization across different state-action pairs.
Learning on one state-action pair will affect the value of all state-action pairs that share similar representations with this pair.
When the embedding network assigns similar representations to similar state-action pairs (e.g., sharing similar distribution over returns), the update for one state-action pair is representative for the updates for other similar state-action pairs,
which improves sample efficiency.
However, such generalization may not be achieved by the base RL algorithm since, when trained by the algorithm with only a return-based loss,
similar state-action pairs
may have similar Q values but very different representations.
{One may argue that, we can adopt several temporal difference updates to propagate the values for state-action pairs with same sampled return, and finally all such pairs are assigned with similar Q values.}
However, since we adopt a learning algorithm with
{bootstrapping/temporal difference learning}
and frozen target network in deep RL, it could take longer time to propagate the value across different state-action pairs, compared with direct generalization over state-action pairs with contrastive learning.
}

\iffalse
\textcolor{red}{
We believe our auxiliary task may boost the learning, due to better return-induced representations that facilitate generalization across different state-action pairs.
A temporal difference (TD) update on one state-action pairs $(s,a)$, will affect the value of all state-actions pairs that share similar representations with
$(s,a)$,
which is induced by return-based losses.
Such generalization may not be achievable in basic RL algorithm, since two state-action pairs $(s, a)$ $(s', a')$ with the same return distribution may have similar Q value but very different representations at the same time.
One may argue that, we can adopt several TD updates to propagate value for pairs with same sampled return. However, since in deep RL we adopt a TD learning algorithm with frozen target network, it could takes longer time to propagate value across different state-action pairs, in contrast to shorter time to generalize over state-action pairs with similar representation.
}
\fi

{
Meanwhile, since we construct auxiliary tasks based on return, which is a very different structure from image augmentation or temporal structure, our method could be combined with existing methods to achieve further improvement. }

\subsection{Network Structure for State-action Embedding}

In our algorithm, the auxiliary task is based on the state-action embedding, instead of the state embedding that is frequently used in the previous work \cite[e.g., ][]{srinivas2020curl}.
To facilitate our algorithm, we design two new structures for Atari 2600 Games~(discrete action) and DMControl Suite~(continuous action) respectively.
We show the structure in Figure \ref{fig:nn_structure}.
For Atari, we learn an action embedding for each action and use the element-wise product of the state embedding and action embedding as the state-action embedding.
For DMControl, the action embedding is a real-valued vector and we use the concatenation of the action embedding and the state embedding.

\begin{figure}[htbp]
    \centering
    \includegraphics[width=0.9\textwidth]{figure4.png}
    \caption{
\textcolor{black}{
The network structure of RCRL for Atari (left) and DMControl Suite (right).
}
    }
    \label{fig:nn_structure}
\end{figure}

\subsection{Proof of Corollary \ref{corollary:main}}

Recall that Z-learning aims to solve the following optimization problem:

\begin{equation}
\label{eq:optimization_app}
    \min_{
    \phi \in \Phi_N, w\in\mathcal{W}_N
    } \mathcal{L}(\phi, w;\mathcal{D}) := \mathbb{E}_{(x_1, x_2, y)\sim \mathcal{D}} \left[ (w(\phi(x_1), \phi(x_2)) - y)^2 \right],
\end{equation}
which can also be regarded as finding a compound predictor $f(\cdot, \cdot) := w(\phi(\cdot), \phi(\cdot))$ over the function class $\mathcal{F}_N:=\{(x_1, x_2)\to w(\phi(x_1), \phi(x_2)): w\in\mathcal{W}_N, \phi\in\Phi_N\}$.

For the first step, it is helpful to find out the Bayes optimal predictor $f^*$ when size of the dataset is infinite.
We notice
the fact that
the Bayes optimal predictor for a square loss is the conditional mean, i.e., given a distribution $D$, the Bayes optimal predictor $f^* = arg \min_f \mathbb{E}_{(x,y)\sim D}[(f(x)-y)^2]$ satisfies $f^*(x')=\mathbb{E}_{(x,y)\sim D}[y\mid x=x']$.
Using this property, we can obtain the Bayes optimal predictor over all the functions $\{\mathcal{X}\times \mathcal{X} \to [0,1]\}$ for our contrastive loss:
\begin{equation}
\label{eq:bayes_optimal}
\begin{aligned}
f^*(x_1, x_2) & = \mathbb{E}_{(x'_1, x'_2, y)\sim D} [y\mid x'_1=x_1, x'_2=x_2] \\
& = \mathbb{E}_{R_1 \sim Z^\pi(x_1), R_2 \sim Z^\pi(x_2)} \mathbb{I}[b(R_1) \neq b(R_2)] \\
&= 1- \mathbb{Z}^\pi(x_1)^T \mathbb{Z}^\pi(x_2),
\end{aligned}
\end{equation}
where we use $D$ to denote the distribution
from which each tuple in the dataset $\mathcal{D}$ is drawn.

To establish the theorem, we require that such an optimal predictor $f^*$ to be in the function class
$\mathcal{F}_N$.
Following a similar argument to Proposition 10 in \citet{misra2019kinematic}, it is not hard to show that using $N > N_{\pi,K}$ is sufficient for this realizability condition to hold.

{\textbf{Corollary 4.1.1.}\ \textit{The encoder $\hat{\phi}$ returned by Algorithm \ref{algo:z-learning} with $n\to\infty$ is a $Z^\pi$-irrelevance,
i.e.,
for any $x_1, x_2\in\mathcal{X}$, $\mathbb{Z}^\pi(x_1)
= \mathbb{Z}^\pi(x_2)$ if $\hat{\phi}(x_1)=\hat{\phi}(x_2)$.}
}

\begin{proof}[Proof of Corollary \ref{corollary:main}]
Considering the asymptotic case (i.e., $n\to\infty$), we have $\hat{f} = f^*$ where $\hat{f}(\cdot, \cdot):=\hat{w}(\hat{\phi}(\cdot), \hat{\phi}(\cdot))$ and $\hat{w}$ and $\hat{\phi}$ is returned by Algorithm \ref{algo:z-learning}. If $\hat{\phi}(x_1) = \hat{\phi}(x_2)$, we have for any $x \in \mathcal{X}$,
\begin{equation*}
\begin{aligned}
    1 - \mathbb{Z}^\pi(x_1)^T Z^\pi(x) = f^*(x_1, x) = \hat{f}(x_1, x) = \hat{f}(x_2, x) = f^*(x_2, x) = 1 - \mathbb{Z}^\pi(x_2)^T \mathbb{Z}^\pi(x).
\end{aligned}
\end{equation*}
We obtain $\mathbb{Z}^\pi(x_1) = \mathbb{Z}^\pi(x_2)$ by letting $x=x_1$ or $x=x_2$.
\end{proof}

\subsection{Proof of Theorem \ref{theorem:main}}

\textcolor{black}{
\textbf{Theorem 4.1.}\ \textit{Given the encoder $\hat{\phi}$ returned by Algorithm \ref{algo:z-learning}, the following inequality holds with probability $1-\delta$ and for any $x'\in\mathcal{X}$:}
\begin{equation}
\begin{aligned}
\mathbb{E}_{x_1 \sim d, x_2\sim d} & \Big[ \mathbb{I}[\hat{\phi}(x_1)=\hat{\phi}(x_2)] \, \Big|
 \mathbb{Z}^\pi(x')^T \left( \mathbb{Z}^\pi(x_1) - \mathbb{Z}^\pi(x_2) \right)
\Big| \Big] \\
& \quad \quad \le \sqrt{\frac{8N}{n} \Big( 3 + 4N^2 \ln n + 4 \ln |\Phi_N| + 4 \ln (\frac{2}{\delta}) \Big) },
\end{aligned}
\end{equation}
\textit{where $|\Phi_N|$ is the cardinality of encoder function class and $n$ is the size of the dataset.}
}

\textcolor{black}{The theorem shows that
1) whenever $\hat{\phi}$ maps two state-actions $x_1, x_2$ to the same abstraction, $\mathbb{Z}^\pi(x_1) \approx \mathbb{Z}^\pi(x_2)$ up to an error proportional to $1/\sqrt{n}$ (ignoring the logarithm factor), and
2) when the difference of the return distributions (e.g., $\mathbb{Z}^\pi(x_1) - \mathbb{Z}^\pi(x_2)$) is large, the chance that two state-action pairs (e.g., $x_1$ and $x_2$) are assigned with the same encoding is small.
However, since state-action pairs are sampled i.i.d. from the distribution $d$, the bound holds in an average sense instead of the worse-case sense.
}

\textcolor{black}{
The overview of the proof is as follows:
Note that the theorem builds connection between the optimization problem defined in \eqref{eq:optimization_app} and the learned encoder $\hat{\phi}$.
To prove the theorem, we first find out the minimizer $f^*$ of the optimization problem when the number of samples is infinite (which was calculated in \eqref{eq:bayes_optimal} in the previous subsection).
Then, we bound the difference between the learned predictor $\hat{f}$ and $f^*$ in Lemma \ref{lemma:concentration} (which utilizes Lemma \ref{lemma:pointwise_covering_number}) when the number of samples is finite.
At last, utilizing the compound structure of $\hat{f}$, we can relate it to the encoder $\hat{\phi}$.
Specifically, given $x_1, x_2 \in \mathcal{X}$ with $\hat{\phi}(x_1) = \hat{\phi}(x_2)$, we have $\hat{f}(x_1, x') = \hat{f}(x_2, x'), \forall x'\in\mathcal{X}$.
}

\begin{definition}[Pointwise covering number]
Given a function class $\mathcal{G}: \mathcal{X} \to \mathbb{R}$, the pointwise covering number at scale $\epsilon$, $\mathcal{N}(\mathcal{G}, \epsilon)$, is the size of the set $V:\mathcal{X} \to \mathbb{R}$ such that $\forall g \in \mathcal{G}, \exists v\in V: \sup_{x\in\mathcal{X}} |g(x) - v(x)| < \epsilon$.
\end{definition}

\begin{lemma}
\label{lemma:pointwise_covering_number}
The logarithm of the pointwise covering number of the function class $\mathcal{F}_N$ satisfies
$\ln \mathcal{N}(\mathcal{F}_N, \epsilon) \le N^2 \ln (\frac{1}{\epsilon}) + \ln |\Phi_N|$.
\end{lemma}

\begin{proof}
Recall that $\mathcal{F}_N:=\{(x_1, x_2)\to w(\phi(x_1), \phi(x_2)): w\in\mathcal{W}_N, \phi\in\Phi_N\}$, where $\mathcal{W}_N:=\{[N]\times[N]\to[0,1]\}$ and $\Phi_N:=\{ \mathcal{X}\to[N]\}$.
Given $\epsilon > 0$, we discretize $[0,1]$ to $Y := \{\epsilon, \cdots, \lceil \frac{1}{\epsilon} \rceil \epsilon \}$.
Then, we define $W_N : =\{ [N]\times [N] \to Y \}$ and it is easy to see that $W_N$ is a covering of $\mathcal{W}_N$ with $|W_N| \le (\frac{1}{\epsilon})^{N^2}$.
Next, we observe that $F_N := \{(x_1, x_2)\to w(\phi(x_1), \phi(x_2)): w\in W_N, \phi\in\Phi_N\}$ is a covering of $\mathcal{F}_N$ and $|F_N| = |\Phi_N| \, |W_N|$.
We complete the proof by combining the above results.
\end{proof}

\begin{proposition}[Proposition 12 in \citet{misra2019kinematic}]
\label{proposition:concentration}
Consider a function class $\mathcal{G}: \mathcal{X} \to \mathbb{R}$ and $n$ samples $\{(x_i, y_i)\}_{i=1}^n$ drawn from $D$, where $x_i \in \mathcal{X}$ and $y_i \in [0,1]$.
The Bayes optimal function is
$g^* = \text{arg} \min_{g\in\mathcal{G}} \mathbb{E}_{(x,y)\sim D} \left[ (g(x) - y)^2 \right]$ and the empirical risk minimizer is
$\hat{g} = \text{arg} \min_{g\in\mathcal{G}} \frac{1}{n} \sum_{i=1}^n \left[ (g(x_i) - y_i)^2 \right].$
With probability at least $1-\delta$ for a fixed $\delta \in (0,1)$,
$$
\mathcal{E}_{(x,y)\sim D} \left[ (\hat{g}(x) - g^*(x))^2 \right] \le \inf_{\epsilon > 0} \left[
6\epsilon + \dfrac{8 \ln (2\mathcal{N}(\mathcal{G}, \epsilon) / \delta) }{n}
\right],
$$
\end{proposition}

\begin{lemma}
\label{lemma:concentration}
Given the empirical risk minimizer $\hat{f}$ in Algorithm \ref{algo:z-learning}, we have
\begin{equation}
\mathbb{E}_{(x_1, x_2) \sim D} \left[ (\hat{f}(x_1, x_2) - f^*(x_1, x_2))^2 \right] \le \Delta_\text{reg} =
\frac{6}{n} + \frac{8}{n} \left(
N^2 \ln n + \ln |\Phi_N| + \ln (\frac{2}{\delta})
\right),
\end{equation}
\end{lemma}
with probability at least $1-\delta$.

\begin{proof}
This directly follows from the combination of Lemma \ref{lemma:pointwise_covering_number} and Proposition \ref{proposition:concentration} by letting $\epsilon = \frac{1}{n}$.
\end{proof}

\begin{proof}[Proof of Theorem \ref{theorem:main}]
First, we
denote $\mathcal{E}_i := \mathbb{I}[\hat{\phi}(x_1) = i = \hat{\phi}(x_2)]$ and $b_i := \mathbb{E}_{x\sim d}\left[ \mathbb{I}[\hat{\phi}(x)=i] \right]$ to be the prior probability over the $i$-th abstraction.
Then, for all $x'\in\mathcal{X}$, we have
\begin{equation*}
\begin{aligned}
& \mathbb{E}_{x_1 \sim d, x_2 \sim d} \left[
\mathbb{I}[\hat{\phi}(x_1) = i = \hat{\phi}(x_2)] \,
\big| \mathbb{Z}^\pi(x')^T \left( Z^\pi(x_1) - Z^\pi(x_2) \right) \big|
\right] \\
= &
\mathbb{E}_{x_1 \sim d, x_2 \sim d} \left[
\mathcal{E}_i \,
|f^*(x_1, x') - f^*(x_2, x')|
\right] \\
\le &
\mathbb{E}_{x_1 \sim d, x_2 \sim d} \left[
\mathcal{E}_i \,
\left(
|f^*(x_1, x') - \hat{f}(x_1, x')| +
|f^*(x_1, x') - \hat{f}(x_2, x')|
\right)
\right] \\
= &
\mathbb{E}_{x_1 \sim d, x_2 \sim d} \left[
\mathcal{E}_i \,
\left(
|f^*(x_1, x') - \hat{f}(x_1, x')| +
|f^*(x_1, x') - \hat{f}(x_1, x')|
\right)
\right] \\
\le &
2 \mathbb{E}_{x_1 \sim d} \left[
\mathbb{I} [\hat{\phi}(x_1) = i] \,
\big| f^*(x_1, x') - \hat{f}(x_1, x') \big|
\right] \\
\le &
2 \sqrt{b_i \Delta_\text{reg}}. \\
\end{aligned}
\end{equation*}
The first step is obtained using Equation (\ref{eq:bayes_optimal}).
The second step is the triangle inequality.
The third step uses the fact that under event $\mathcal{E}_i$, we have $\hat{\phi}(x_1) = \hat{\phi}(x_2)$ and therefore $\hat{f}(x_1, x') = \hat{f}(x_2, x')$.
In the fourth step, we drop the dependence on the other variable, and marginalize over $D_\text{coup}$.
The last step is the Cauchy-Schwarz inequality.

At last, we complete the proof by summing over all the $i\in[N]$ and using the fact $\sum_{i=1}^N \sqrt{b_i} \le \sqrt{N}$ and Lemma \ref{lemma:concentration}.
\end{proof}
\subsection{State abstraction}

\textbf{Definitions of abstractions.}
Aggregating states or state abstraction \citep{li2006towards} can increase sample efficiency in RL (empirically demonstrated by \cite{abel2016near}).
There are different levels of abstractions,
such as homomorphism \citep{ravindran2003smdp,ravindran2004algebraic}, bisimulation \citep{dean1997model,givan2003equivalence}, $Q^\pi$ irrelevance, $Q^*$ irrelevance, $a^*$-irrelevance, and $\pi^*$ irrelevance \citep{li2006towards}, arranged from finest to coarsest.

\textbf{Properties of abstractions.}
People investigate approximate form of different abstractions and how they can be used to bound optimality (i.e., whether abstract MDP permits an optimal solution) \citep{abel2016near,jiang2018notes}.
People also investigate whether different forms of abstractions secure Q-learning:
It is shown that $Q^*$ irrelevance (with a gap assumption) secures Q-learning \citep{li2009unifying,abel2016near}, but Q-learning may break down for $\pi^*$ irrelevance \citep{jong2005state}.

\textbf{Finer or coarser.}
Although finer abstractions lead to lower approximation error, coarser abstractions result in lower estimation error \citep{jiang2015abstraction}, which may be more efficient
\footnote{It is shown that the MSE between $V^*$ and the value of the optimal policy on the abstracted MDP can be written as an addition of an approximation error and an estimation error (since the abstraction is estimated based on a dataset).}.
Moreover, it is shown that finding an bisimulation (even if it exists) requires a number of samples that is exponential in horizon (see Proposition B.1 in \cite{modi2020sample}).

\textbf{From abstraction to metric/representation.}
Similar to bisimulation metric \citep{ferns2004metrics,ferns2011bisimulation} that is proposed to associate bisimulation,
we propose to a $Q^\pi$ irrelevance metric related to the coarser form of abstraction, $Q^\pi$ irrelevance.

\textbf{State-action abstraction.}
Few authors investigate state-action abstractions: homomorphism \cite{ravindran2003smdp} deals with action aggregation and permutation. \cite{maillard2020upper} and \cite{biza2018online} define a state-action based homomorphism.
Here, we study a new form of state-action abstraction (which is similar to $Q^\pi$ irrelevance).

\textbf{Abstraction algorithms.}
In terms of how to obtain abstractions, authors focus on learning a representation associated with bisimulation \citep{ruan2015representation,gelada2019deepmdp,zhang2020learning,van2020mdp}, homomorphism \citep{biza2018online} or $\phi_{a^*}$-irrelevance \citep{jong2005state}.
Particularly, \cite{ortner2013adaptive} use a UCRL algorithm \citep{auer2009near} to quantify the similarity between two states in terms of reward and transition, and further aggregate the states.
DeepMDP \citep{gelada2019deepmdp} and DBC \citep{zhang2020learning} include experiments on complex tasks such as Atari or DMControl and achieve good performance.

\subsection{Auxiliary task}

In reinforcement learning, the auxiliary-task first proposed by \cite{sutton2011horde} can improve sample efficiency.

\textbf{Model-based setting/world model.}
In the model-based settings, world models can be used as auxiliary tasks and lead to better performance, such as CRAR \citep{franccois2019combined}, Dreamer \citep{hafner2019dream} and PlaNet \citep{hafner2019learning}.
These papers learn latent representations of the state end-to-end which can be regarded as aggregating the states implicitly.
However, such methods are unstable to train and relies on regularizations to converge.

\textbf{Model-free setting.}
In the model-free settings, many algorithms construct various auxiliary tasks to improve performance, such as predicting the future states \citep{guo2020bootstrap,jaderberg2016reinforcement}, recognizing from different transformations to the image-based states \citep{srinivas2020curl}, or combination of different auxiliary objectives that aid representation learning \citep{de2018integrating}.
The most relevant work focus on predicting various value functions:
\cite{veeriah2019discovery} predict value functions with different cumulants, \cite{schaul2015universal} and \cite{borsa2018universal} learn value functions with different goals or rewards.
AVF \citep{bellemare2019geometric} and value improvement path \citep{dabney2020value}
predict value functions under different policies, which are similar to ours.
However, we justify such auxiliary tasks through the lens of state abstraction instead of a geometric view which only works when the values are linear combinations of the representation.
Also, our algorithm is much simpler than AVF with only an auxiliary loss in addition to DQN based algorithms.
PastPolicies and PastMixture in \cite{dabney2020value} use auxiliary tasks based on bootstrapping old Q functions or Q functions learned with different learning rates, which is regression based on moving targets.
We design auxiliary tasks based on classification towards fixed targets and we may expect this could result in better performance.

\section{Related Works}
\subsection{Fundamentals}

\begin{itemize}
\item \textbf{Bisimulation}: Equivalence notions and model minimization in Markov decision processes. Robert Givan, Thomas Dean, and Matthew Greig. Artificial Intelligence. 2003.

\item \textbf{Approximate Bisimulation}: Model reduction techniques for computing approximately optimal solutions for Markov decision processes. Thomas Dean, Robert Givan, and Sonia M.Leach. UAI-97.

\item \textbf{Bisimulation Metrics}: Metrics for finite Markov decision processes. Norm Ferns, Prakash Panangaden, and Doina Precup. In Proceedings of the Twentieth Conference on Uncertainty in Artificial Intelligence (UAI-04). 2004.

\item \textbf{Bisimulation Metrics}: Bisimulation metrics for continuous markov decision processes. Ferns, N., Panangaden, P., and Precup, D. 2011.

\item \textbf{MDP Homomorphisms}: An algebraic approach to abstraction in reinforcement learning. Balaraman Ravindran. Phd Thesis. 2004.

\item \textbf{Approximate MDP Homomorphisms}: Approximate Homomorphisms: A Framework for Non-Exact Minimization in Markov Decision Processes. Balaraman Ravindran and Andrew G. Barto. In International Conference on Knowledge Based Computer Systems. 2004.

\item \textbf{MDP Homomorphisms Metrics}: Bounding Performance Loss in Approximate MDP Homomorphisms. NIPS 2008.

\item \textbf{Towards a Unified Theory of State Abstraction for MDPs}. Lihong li, et al. 2006.

\end{itemize}

\subsection{Advances}

\begin{itemize}
    \item \textbf{Value Function Geometry}: Proto-value functions A laplacian framework for learning representation and control in markov decision processes. Journal of Machine Learning Research. Mahadevan, S. and Maggioni, M. 2007.
    \item \textbf{Value Function Geometry}:
    A geometric perspective on optimal representations for reinforcement learning. Bellemare, M. G., et al. ICML 2019.
    \item \textbf{Value Function Geometry}:
    The Value-Improvement Path Towards Better Representations for Reinforcement Learning. Deepmind. arXiv 2020.
    \item \textbf{Auxiliary Loss}:
    Reinforcement learning with unsupervised auxiliary tasks. Jaderberg, M., Mnih, V., Czarnecki, W. M., Schaul, T., Leibo, J. Z., Silver, D., and Kavukcuoglu, K. In International Conference on Learning Representations (ICLR), 2017.
    \item \textbf{Auxiliary Loss}:
    DeepMDP: Learning Continuous Latent Space Models for Representation Learning. ICML 2019.
    \item \textbf{Auxiliary Loss}:
    Learning Invariant Representations for Reinforcement Learning without Reconstruction. Zhang A, McAllister R, Calandra R, et al. arXiv preprint arXiv:2006.10742, 2020.
    \item \textbf{MDP Homomorphisms Metrics Loss}:
    Plannable Approximations to MDP Homomorphisms: Equivariance under Actions. Elise van der Pol, Thomas Kipf, Frans A. Oliehoek, Max Welling. AAMAS 2020.
    \item \textbf{Alternating SRL and RL}: Integrating State Representation Learning into Deep Reinforcement Learning. ICRA 2018.

\end{itemize}
\subsection{Auxiliary task}
In reinforcement learning, the auxiliary task can be used for both the model-based setting and the model-free setting.
In the model-based settings, world models can be used as auxiliary tasks and lead to better performance, such as CRAR~\citep{franccois2019combined},
Dreamer~\citep{hafner2019dream}, and PlaNet~\citep{hafner2019learning}.
Due to the complex components (e.g., the latent transition or reward module) in the world model, such methods are empirically unstable to train and relies on different regularizations to converge.
In the model-free settings, many algorithms construct various auxiliary tasks to improve performance, such as predicting the future \citep{jaderberg2016reinforcement,shelhamer2016loss,guo2020bootstrap,lee2020predictive,mazoure2020deep},
learning value functions with different rewards or under different policies~\citep{veeriah2019discovery,schaul2015universal,borsa2018universal,bellemare2019geometric,dabney2020value}, learning from many-goals~\citep{veeriah2018many},
or the combination of different auxiliary objectives
~\citep{de2018integrating}.
Moreover, auxiliary tasks can be designed based on the prior knowledge about the environment
{\citep{mirowski2016learning,shen2020auxiliary,van2020mdp}}
or the raw state representation
~\citep{srinivas2020curl}.
\citet{hessel2019multi} also apply auxiliary task to the multi-task RL setting.

Contrastive learning has seen dramatic progress recently,
and been introduced to learn state representation~\citep{oord2018representation,sermanet2018time,dwibedi2018learning,aytar2018playing,anand2019unsupervised,srinivas2020curl}.
Temporal structure~\citep{sermanet2018time, aytar2018playing} and local spatial structure~\citep{anand2019unsupervised} has been leveraged for state representation learning via contrastive losses.
CPC~\citep{oord2018representation} and CURL~\citep{srinivas2020curl} adopt a contrastive auxiliary tasks to accelerate representation learning and speed up main RL tasks,
{by leveraging the temporal structure and image augmentation respectively. }
{
To the best of our knowledge, we are the first to leverage return to construct a contrastive auxiliary task for speeding up the main RL task.}

\subsection{Abstraction}

{
State abstraction (or state aggregation) aggregates states by ignoring irrelevant state information. By reducing state space, state abstraction can enable efficient policy learning.
Different types of abstraction are proposed in literature, ranging from fine-grained to coarse-grained abstraction, each reducing state space to a different extent.
Bisimulation or model irrelevance \citep{dean1997model,givan2003equivalence} define state abstraction under which both transition and reward function are kept invariant.
By contrast, other types of state abstraction that are coarser than bisimulation such as $Q^\pi$ irrelevance or $Q^*$ irrelevance \citep[see e.g.,][]{li2006towards}, which keep the Q function invariant under any policy $\pi$ or the optimal policy respectively.
There are also some works on state-action abstractions, e.g., MDP homomorphism \citep{ravindran2003smdp,ravindran2004algebraic}
\textcolor{black}{and approximate MDP homomorphism \citep{ravindran2004approximate,taylor2009bounding}}
, which are similar to bisimulation in keeping reward and transition invariant, but extending bisimulation from state abstraction to state-action abstraction.
}

{
In this paper, we consider a new form of state-action abstraction $Z^\pi$-irrelevance, which
aggregates state-action pairs with the same return distribution and
is coarser than bisimulation or homomorphism which are frequently used as auxiliary tasks \textcolor{black}{\citep[e.g., ][]{biza2018online,gelada2019deepmdp,zhang2020learning}}.}
However, it is worth noting that $Z^\pi$-irrelevance is only used to build the theoretical foundation of our algorithm, and show that our proposed auxiliary task is well-aligned with the main RL task. Representation learning in deep RL is in general very different from aggregating states in tabular case, though the latter may build nice theoretical foundation for the former. Here we focus on how to design auxiliary tasks to accelerate representation learning using contrastive learning techniques, and we propose a novel return-based contrastive method based on our proposed $Z^\pi$-irrelevance abstraction.

\iffalse
Sample complexity is polynomial in the size of the state space for tabular RL algorithms \citep{jiang2018open}.
When the size of the state space becomes large, as will be the case in many real world problems, the agent needs to generalize what is learned about one state to other states.
One of the popular generalization schemes is abstraction \citep[see][]{jiang2018notes}, and different types of abstraction are proposed in literature.
Bisimulation or model irrelevance \citep{dean1997model,givan2003equivalence} is the finest abstraction in the abstraction hierarchy and is frequently adapted as auxiliary tasks \citep{ruan2015representation,gelada2019deepmdp,zhang2020learning}.
By contrast, other types of state abstraction that are coarser than bisimulation such as $Q^\pi$ irrelevance or $Q^*$ irrelevance \citep[see e.g.,][]{li2006towards} are less frequently used as auxiliary tasks.
Nevertheless, evidences exist that a coarser abstraction may be a better choice for using as an auxiliary task:
While finer abstractions lead to lower approximation error, coarser abstractions result in lower estimation error \citep{jiang2015abstraction}.
Compared with approximation error, estimation error in learning a state abstraction seems to be an even more important problem.
It is shown that finding an bisimulation (even if it exists) requires a number of samples that is exponential in horizon (see Proposition B.1 in \cite{modi2020sample}).

There are also some works on state-action abstractions, i.e., MDP homomorphism \citep{ravindran2003smdp,ravindran2004algebraic, van2020mdp,biza2018online},
which is similar to bisimulation in keeping reward and transition invariant under abstration, but extending bisimulation from state abstraction to state-action abstraction.

In this paper, we consider a new form of state-action abstraction, which is coarser than existing ones and aggregates
state-action pairs
with the same return distribution.
Combined with contrastive learning techniques, we propose an algorithm called Z-learning to learn this state-action abstraction in a deep-learning-friendly way, which is later adapted to RCRL.
\fi

\subsection{Abstraction and its properties}

This part follows a standard infinite horizon discounted MDP setting/notations. In addition, we use $x\in\mathcal{X}$ to denote the state-actions, i.e., $\mathcal{X} = \mathcal{S} \times \mathcal{A}$. Let $n=|\mathcal{X}|$.

\begin{definition}[$Q$-irrelevance abstraction]
Given a set of policies $\Pi$ and any state-actions $x,\bar{x} \in \mathcal{X}$,
a $Q$-irrelevance abstraction $\phi_{Q,\Pi}$ is such that $\phi_{Q,\Pi}(x) = \phi_{Q,\Pi}(\bar{x})$ implies $Q^\pi(x)=Q^\pi(\bar{x})$ for any policy $\pi \in \Pi$.
We also define an exact $Q$-irrelevance abstraction $\varphi_{Q,\Pi}$ such that $\varphi_{Q,\Pi}(x) = \varphi_{Q,\Pi}(\bar{x})$ if and only if $Q^\pi(x)=Q^\pi(\bar{x})$ for any policy $\pi \in \Pi$.
\end{definition}

\begin{definition}[\cite{li2006towards}]
Given two abstractions $\phi_1$ and $\phi_2$, we say $\phi_1$ is finer than $\phi_2$, denoted $\phi_1 \succeq \phi_2$, if and only if $\forall x_1, x_2 \in \mathcal{X}, \phi_1(x_1) = \phi_1(x_2)$ implies $\phi_2(x_1) = \phi_2(x_2)$.
If, in addition, $\phi_1 \neq \phi_2$, then $\phi_1$ is strictly finer than $\phi_2$, denoted $\phi_1 \succ \phi_2$.
We also may say $\phi_2$ is (strictly) coarser than $\phi_1$, denoted $\phi_2 \preceq \phi_1$ ($\phi_2 \prec \phi_1$).
\end{definition}

\begin{remark}
$\phi_{Q,\Pi}$ is the set of abstractions that are finer than $\varphi_{Q,\Pi}$.
\end{remark}

\begin{definition}[Stationary policies]
The set of all stationary policies is defined as $\Pi_0 = \{\pi \in [0,1]^{\mathcal{S} \times \mathcal{A}} : \sum_{a\in\mathcal{A}}\pi(s,a) = 1\}$. The set of all deterministic stationary policies is defined as
$\Pi_1 = \{\pi \in \{0,1\}^{\mathcal{S} \times \mathcal{A}} : \sum_{a\in\mathcal{A}}\pi(s,a) = 1\}$.
\end{definition}

\begin{lemma}
\label{lemma:detpolicy}
Given $x, \bar{x} \in \mathcal{X}$, if $Q^{\pi_0}(x) \neq Q^{\pi_0}(\bar{x})$ for some ${\pi_0} \in \Pi_0$, there exists $\pi_1 \in \Pi_1$ such that $Q^{\pi_1}(x) \neq Q^{\pi_1}(\bar{x})$.
\end{lemma}

\begin{proof}
Assume $Q^{\pi_0}(x) > Q^{\pi_0}(\bar{x})$ or we can swap $x$ and $\bar{x}$. Let $f(\pi) = Q^\pi(x) - Q^\pi(\bar{x}) = \sum_{s'} \left( p(s'|x) - p(s'|\bar{x}) \right) V^\pi(s') = \delta^T V^\pi$,
where $\delta \in \mathbb{R}^S$ is a vector, the $s'$-th element of which is $ p(s'|x) - p(s'|\bar{x})$.
According to Lemma 1 in \cite{bellemare2019geometric}, there exists a deterministic policy $\pi_1$ that maximizes $f(\pi)$.
Since there exists a (possible stochastic) policy $\pi_0$ such that $f(\pi_0) > 0$, $f(\pi_1) \ge f(\pi_0) > 0$.
Therefore, $Q^{\pi_1}(x) \neq Q^{\pi_1}(\bar{x})$.
\end{proof}

\begin{theorem}
$\varphi_{Q,\Pi_0} = \varphi_{Q,\Pi_1}$.
\end{theorem}

\begin{proof}
Since $\Pi_1 \subset \Pi_0$, $\varphi_{Q,\Pi_0} \succeq \varphi_{Q,\Pi_1}$.
For any $x,\bar{x} \in \mathcal{X}$ such that $\varphi_{Q,\Pi_1}(x) = \varphi_{Q,\Pi_1}(\bar{x})$, $Q^{\pi_1}(x) = Q^{\pi_1}(\bar{x}), \forall \pi_1 \in \Pi_1$.
Using Lemma \ref{lemma:detpolicy}, this implies that $Q^{\pi_0}(x) = Q^{\pi_0}(\bar{x}), \forall \pi_0 \in \Pi_0$ and therefore $\varphi_{Q,\Pi_0}(x) = \varphi_{Q,\Pi_0}(\bar{x})$, i.e., $\varphi_{Q,\Pi_0} \preceq \varphi_{Q,\Pi_1}$.
\end{proof}

Next we define the corresponding metric and representation of an abstraction.

\begin{definition}[Abstraction metric, \cite{ferns2004metrics}]
Let $\phi$ be an abstraction and $\mathcal{M}$ be the set of all semimetrics on $\mathcal{X}$ that assign distances of at most 1.
Note that every semimetric $d$ induces an abstraction $\phi_d$, obtained by equating points assigned distance zero by $d$.
We say that $d\in\mathcal{M}$ is an abstraction metric of $\phi$ if $\phi_d=\phi$.
\end{definition}

\begin{remark}
$d(x,\bar{x}) := \mathbb{I}\{\phi(x)\neq \phi(\bar{x})\}$ that assigns distance $0$ if $\phi(x) = \phi(\bar{x})$, otherwise assigns distance $1$, is a simple abstraction metric.
However, this processes no more distinguishing power than $\phi$ and we would like the metric distances vary smoothly.
\end{remark}

\begin{definition}[Abstraction representation]
We say that a state-action representation $\psi: \mathcal{X} \to \mathbb{R}^m$ is an abstraction representation of $\phi$ if: $\psi(x) = \psi(\bar{x})$ if and only if $\phi(x) = \phi(\bar{x})$, where $m$ is the dimension of the representation.
\end{definition}

\begin{theorem}
Given an abstraction representation $\psi$ of $\phi_{Q, \Pi_0}$ and any policy  $\pi \in \Pi_0$, there exists a function ${Q}: \mathbb{R}^m \to \mathbb{R}$ such that ${Q}({\psi}(x)) = Q^\pi(x), \forall x \in \mathcal{X}$.
Specifically, there exists ${Q}$ such that ${Q}({\psi}(x)) = Q^*(x), \forall x \in \mathcal{X}$.
\end{theorem}

\begin{proof}
Given a policy $\pi \in \Pi_0$, we can construct a ${Q}$ such that ${Q}({\psi}(x)) = Q^\pi(x), \forall x \in \mathcal{X}$ in the following way:
For all $x\in\mathcal{X}$, one by one, we assign $Q(y) \leftarrow Q^\pi(x)$, where $y = \psi(x)$.
Then, we assign $Q(y)$ with arbitrary values if $y$ does not equal to the representation of any state-actions.
For any $x\in\mathcal{X}$ with $y=\psi(x)$, $Q(y)$ equals either to $Q^\pi(x)$ or $Q^\pi(\bar{x})$ for some other $\bar{x}$ that satisfies $\psi(\bar{x}) = \psi(x)$. However, in the later case, $Q^\pi(x) = Q^\pi(\bar{x})$ since $\psi$ is a state-action representation corresponding to $\phi_{Q, \Pi_0}$.
Therefore, ${Q}({\psi}(x)) = Q^\pi(x), \forall x \in \mathcal{X}$.
This also applies to the optimal policy $\pi^*$.
\end{proof}

The theorem shows that such a state-action representation is fine enough to represent any Q functions. However, to make it complete, we should either 1) prove that any abstraction coarser than $\varphi_{Q, \Pi_0}$ may not preserve Q functions, or at least 2) $\varphi_{Q, \Pi_0}$ is coarser than homomorphism.

\subsection{Algorithm to obtain abstraction}

\begin{assumption}
\label{assumption:detMDP}
We assume a standard tabular case, where the states and actions are discretized.
We also assume the MDP has deterministic dynamics and a fixed initial state $s_0$.
\end{assumption}

\begin{remark}
Tabular case is a standard setting for MDP analysis.
Deterministic dynamics and fixed initial state correspond to "no sticky action" and "no random no-op actions at the start" respectively in Atari games.
Determinism is essential to determine the Q values with a single rollout.
\end{remark}

\begin{algorithm}[H]
\label{algo:exact}
\SetAlgoLined
    \KwIn {A set of deterministic policies $\Pi\subseteq \Pi_1$}
    \KwOut {An abstraction metric of $\varphi_{Q, \Pi}$}
    \ForEach{$\pi$ in $\Pi$}{
        Roll out $\pi$ till end and collect the state-actions and the returns $D_\pi = \{(x_i, R_i)\}$

        $D \leftarrow D \cup \{(x_i, x_j, y): x_i \in \mathcal{X}, x_j \in \mathcal{X}, y = 1 - \mathbb{I}[ (x_i, R_i) \in D_\pi \text{ \& } (x_j, R_j) \in D_\pi \text{ \& } R_i = R_j ]
        \}
        $
    }

    $d \leftarrow arg \min_{d} L(d;D)$, where \begin{equation}
    \label{eq:loss_d}
    L(d;D) = - \mathbb{E}_{(x_i, x_j, y) \sim D}
    \left[
    y \log d(x_i, x_j) + (1-y) \log ( 1- d(x_i, x_j) )

    \right]
    \end{equation}

    \Return{d}
\caption{Learn an abstraction representation of a Q-irrelevance abstraction}
\end{algorithm}

\begin{remark}[Number of steps in one rollout]
Since we assume a discounted setting with infinite horizon, the number of steps in one rollout may be infinite.
However, we assume there exists an integer $T>0$ and the agent will enter an absorbing state $s_T$ within $T$ steps and receive no more reward from $s_T$.
Therefore, there will be no more than $T$ state-actions in one rollout.
This is a natural extension of the episodic setting to the infinite horizon setting, and is suitable for Atari games.
\end{remark}

\begin{remark}[Stochastic version]
The size of $D$ is $n^2|\Pi|$.
However, we can perform stochastic gradient update which avoids storing all the transitions.
\end{remark}

\begin{remark}[Deterministic policies]
Trajectories rolled out by any (possibly stochastic) policy can be considered as being rolled out by a deterministic policy, as long as we choose the same actions from the same states in one trajectory.
\end{remark}

\begin{remark}[Iterating over the state-action space]
We assume that we have access to the state-action space.
\end{remark}

Given a set of deterministic policies $\Pi$, although the abstraction $\varphi_{Q,\Pi}$ may be unique, there are more than one abstraction metrics of $\varphi_{Q,\Pi}$.
We will show that
the following abstraction metric 1) is a metric of $\varphi_{Q,\Pi}$, and 2) is the metric outputted by Algorithm \ref{algo:exact}.

\begin{equation}
d_1(x_1, x_2) = 1 - \dfrac{\sum_{\pi \in \Pi} \mathbb{I}  [ x_1 \in \mathcal{X}_\pi \text{ \& } x_2 \in \mathcal{X}_\pi \text{ \& } Q^\pi(x_1) = Q^\pi(x_2) ]
}{|\Pi|},
\end{equation}
where $\mathcal{X}_\pi$ is the set of all the state-actions that can be visited starting from $s_0$ and following the policy $\pi$.

\begin{theorem}
$d^*$ is an abstraction metric of $\varphi_{Q, \Pi}$.
\end{theorem}

\begin{proof}
We can see that $d_1(x_1, x_2)=0$ $\Leftrightarrow$ $x_1$ and $x_2$ have the same Q values for any policy in $\Pi$ $\Leftrightarrow$ $\varphi_{Q,\Pi}(x_1) = \varphi_{Q,\Pi}(x_2)$.
Also, it can be shown that $d_1$ is
a semimetric that assign distances of at most $1$:
\begin{itemize}
    \item Identity of indiscernibles: $x_1 = x_2 \Rightarrow d_1(x_1, x_2) = 0$;
    \item Symmetry: $d_1(x_1, x_2) = d_1(x_2, x_1)$;
    \item Triangle inequality: $|\Pi|(2 - d_1(x_1, x_2) - d(x_1, x_3)) = \sum_\pi \mathbb{I}[x_1\in\mathcal{X}_\pi \text{ \& } x_2\in\mathcal{X}_\pi \text{ \& } Q^\pi(x_1) = Q^\pi(x_2)] + \mathbb{I}[x_1\in\mathcal{X}_\pi \text{ \& } x_3\in\mathcal{X}_\pi \text{ \& } Q^\pi(x_1) = Q^\pi(x_3)] \le \sum_\pi \mathbb{I}[x_1\in\mathcal{X}_\pi \text{ \& } x_2\in\mathcal{X}_\pi \text{ \& } x_3\in\mathcal{X}_\pi \text{ \& } Q^\pi(x_1) = Q^\pi(x_2) = Q^\pi(x_3)] + |\Pi| \le \sum_\pi \mathbb{I}[x_1\in\mathcal{X}_\pi \text{ \& } x_3\in\mathcal{X}_\pi \text{ \& } Q^\pi(x_1) = Q^\pi(x_3)] + |\Pi| = |\Pi|(2 - d_1(x_1, x_3))$
    \item Boundedness: $0 \le d_1(x_1, x_2) \le 1$;
\end{itemize}
\end{proof}

\begin{remark}
If two state-actions (which may be essentially similar) do not appear in a same trajectory for some policy, $d_1$ will treat them as different state-actions.
This style of aggregation seems too strict.
\end{remark}

\begin{theorem}
Under the Assumption \ref{assumption:detMDP}, Algorithm \ref{algo:exact} returns an abstraction metric that equals to $d_1$.
\end{theorem}

\begin{proof}
Since the environment (i.e., the dynamics of the MDP and the initial state) and the policies are deterministic, the (discounted) return is equal to the Q value.

By letting
\begin{equation*}
\begin{aligned}
\dfrac{\partial L(d;D)}{\partial d(x_i, x_j)}
& \propto \dfrac{\partial}{\partial d(x_i, x_j)} \bigg(
d_1(x_i, x_j) \log d(x_i, x_j) + (1-d_1(x_i, x_j)) \log (1 - d(x_i, x_j))
\bigg)
\\
& = 0,
\end{aligned}
\end{equation*}
we obtain $d(x_i, x_j) = d_1(x_i, x_j)$.

\end{proof}

In practice, instead of learning an abstraction metric, we could learn an abstraction representation $\psi: \mathcal{X} \to \mathbb{R}^m$.
In this case, we could either solve the following optimization problem:

\begin{equation}
\begin{aligned}
    & \max_{\psi} \mathbb{E}_{(x_i, x_j, x_k) \sim D}
    \left[
    \log \dfrac{1-\psi(x_i)^T \psi(x_j)}{1+\psi(x_i)^T \psi(x_j)} + \beta \log (1+\psi(x_i)^T \psi(x_k))
    \right] \\
    & s.t. ||\psi(x)|| = 1, \forall x \in \mathcal{X} \\
\end{aligned}
\end{equation}

or train a DNN to generate the abstraction metric with the abstraction representations as inputs, e.g., $F(\psi(x_1), \psi(x_2)) \in [0,1]$, where $F$ and $\psi$ are DNNs.

\begin{corollary}
Under the Assumption \ref{assumption:detMDP} and given $\Pi = \Pi_1$, Algorithm \ref{algo:exact} returns an abstraction metric of $\varphi_{Q, \Pi_0}$, and the corresponding representation preserves the value functions for any policy.
\end{corollary}

\subsection{Towards practical algorithms}

The condition for aggregation in $d_1$ is too strict and Algorithm \ref{algo:exact} requires to access the full state-action space.
Here, we provide another modified version of algorithm to alleviate these two problems.

\begin{algorithm}[H]
\label{algo:2}
\SetAlgoLined
    \KwIn {A set of deterministic policies $\Pi\subseteq \Pi_1$}
    \KwOut {An abstraction metric of $\varphi_{Q, \Pi}$}
    \ForEach{$\pi$ in $\Pi$}{
        Roll out $\pi$ till end and collect the state-actions and the returns $D_\pi = \{(x_i, R_i)\}$

        $D \leftarrow D \cup \{(x_i, x_j, y): (x_i, R_i) \in D_\pi, (x_j, R_j) \in D_\pi, y=\mathbb{I}[ R_i \neq R_j ]
        \}
        $
    }

    $d \leftarrow arg \min_{d} L(d;D)$, where \begin{equation}
    \label{eq:loss_d}
    L(d;D) = - \mathbb{E}_{(x_i, x_j, y) \sim D}
    \left[
    y \log d(x_i, x_j) + (1-y) \log ( 1- d(x_i, x_j) )

    \right]
    \end{equation}

    \Return{d}
\caption{Learn an abstraction representation of a Q-irrelevance abstraction}
\end{algorithm}

Similarly, we show that this algorithm outputs the following solution:

\begin{equation}
d_2(x_1, x_2) = \dfrac{\sum_{\pi \in \Pi} \mathbb{I}  [ x_1 \in \mathcal{X}_\pi \text{ \& } x_2 \in \mathcal{X}_\pi \text{ \& } Q^\pi(x_1) \neq Q^\pi(x_2) ]
}{\sum_{\pi \in \Pi} \mathbb{I}  [ x_1 \in \mathcal{X}_\pi \text{ \& } x_2 \in \mathcal{X}_\pi ]},
\end{equation}

\begin{theorem}
Under the Assumption \ref{assumption:detMDP}, Algorithm \ref{algo:2} returns $d_2$.
\end{theorem}

\begin{proof}
Similarly, by letting
\begin{equation*}
\begin{aligned}
\dfrac{\partial L(d;D)}{\partial d(x_i, x_j)}
& = \dfrac{\partial}{\partial d(x_i, x_j)} \bigg(
\sum_{\pi\in\Pi} \mathbb{I}[x_i \in \mathcal{X}_\pi \text{ \& } x_j \in \mathcal{X}_\pi] \\
&
\quad \quad \left(
\mathbb{I}[R_i \neq R_j] \log d(x_i, x_j) + (1 - \mathbb{I}[R_i \neq R_j]) (1 - d(x_i, x_j))
\right)
\bigg)
\\
& \propto \dfrac{\partial}{\partial d(x_i, x_j)} \bigg(
d_2(x_i, x_j) \log d(x_i, x_j) + (1-d_2(x_i, x_j)) \log (1 - d(x_i, x_j))
\bigg)
\\
& = 0,
\end{aligned}
\end{equation*}
we obtain $d(x_i, x_j) = d_2(x_i, x_j)$.

\end{proof}

\begin{lemma}[Relationship between $d_1$ and $d_2$]
Given $x_1, x_2 \in \mathcal{X}$, $d_1(x_1, x_2) \ge d_2(x_1, x_2)$.
Specifically, if $d_1(x_1, x_2) = 0$ for some $x_1, x_2 \in \mathcal{X}$, $d_2(x_1, x_2) = 0$;
if $d_2(x_1, x_2) = 1$ for some $x_1, x_2 \in \mathcal{X}$, $d_1(x_1, x_2) = 1$.
\end{lemma}

\begin{remark}[$d_2$ is the lower bound of $d_1$]
If we do not have access to the full state-action space, we construct $D$ with only the encountered state-actions. Then, we will obtain $d_2$ which is the lower bound of the abstraction metric $d_1$. Notice that $d_2$ is not a semimetric since it does not satisfy triangle inequality.
\end{remark}

\begin{remark}[Practical algorithm that incorporates DQN/Rainbow]
Further, this motivates us to use the trajectories in the replay buffer to construct $D$ and use $L(d; D)$ as the loss function to learn the approximate abstraction representation of $\varphi_{Q, \Pi}$ where $\Pi$ is the set of the deterministic policies induced by the transitions in the replay buffer.
Moreover, we use this as the auxiliary task along with a standard RL algorithm.
\end{remark}

redpajama_set_name: RedPajamaArXiv