# DR3: Value-Based Deep Reinforcement
Learning Requires Explicit Regularization

Aviral Kumar1,2 Rishabh Agarwal2,3
Tengyu Ma4 Aaron Courville3 George Tucker2 Sergey Levine1,2
1 UC Berkeley 2 Google Research 3 MILA 4 Stanford University
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Despite overparameterization, deep networks trained via supervised learning
are easy to optimize and exhibit excellent generalization. One hypothesis to
explain this is that overparameterized deep networks enjoy the benefits of
implicit regularization induced by stochastic gradient descent, which favors
parsimonious solutions that generalize well on test inputs. It is reasonable
to surmise that deep reinforcement learning (RL) methods could also benefit
from this effect. In this paper, we discuss how the implicit regularization
effect of SGD seen in supervised learning could in fact be harmful in the
offline deep RL setting, leading to poor generalization and degenerate feature
representations. Our theoretical analysis shows that when existing models of
implicit regularization are applied to temporal difference learning, the
resulting derived regularizer favors degenerate solutions with excessive
“aliasing”, in stark contrast to the supervised learning case. We back up
these findings empirically, showing that feature representations learned by a
deep network value function trained via bootstrapping can indeed become
degenerate, aliasing the representations for state-action pairs that appear on
either side of the Bellman backup. To address this issue, we derive the form
of this implicit regularizer and, inspired by this derivation, propose a
simple and effective _explicit_ regularizer, called DR3, that counteracts the
undesirable effects of this implicit regularizer. When combined with existing
offline RL methods, DR3 substantially improves performance and stability,
alleviating unlearning in Atari 2600 games, D4RL domains and robotic
manipulation from images.

###  1 Introduction

Deep neural networks are overparameterized, with billions of parameters, which
in principle should leave them vulnerable to overfitting. Despite this,
supervised learning with deep networks still learn representations that
generalize well. A widely held consensus is that deep nets find simple
solutions that generalize due to various _implicit_ regularization effects
(Blanc et al., 2020; Woodworth et al., 2020; Arora et al., 2018; Gunasekar et
al., 2017; Wei et al., 2019; Li et al., 2019). We may surmise that using deep
neural nets in reinforcement learning (RL) will work well for the same reason,
learning effective representations that generalize due to such implicit
regularization effects. But is this actually the case for value functions
trained via bootstrapping?

In this paper, we argue that, while implicit regularization leads to effective
representations in supervised deep learning, it may lead to poor learned
representations when training overparameterized deep network value functions.
In order to rule out confounding effects from exploration and non-stationary
data distributions, we focus on the offline RL setting – where deep value
networks must be trained from a static dataset of experience. There is already
evidence that value functions trained via bootstrapping learn poor
representations: value functions trained with offline deep RL eventually
degrade in performance (Agarwal et al., 2020; Kumar et al., 2021) and this
degradation is correlated with the emergence of low-rank features in the value
network (Kumar et al., 2021). Our goal is to understand the underlying cause
of the emergence of poor representations during bootstrapping and develop a
potential solution. Building on the theoretical framework developed by Blanc
et al. (2020); Damian et al. (2021), we characterize the implicit regularizer
that arises when training deep value functions with TD learning. The form of
this implicit regularizer implies that TD-learning would co-adapt feature
representations at state-action tuples that appear on either side of a Bellman
backup.

We show that this theoretically predicted aliasing phenomenon manifests in
practice as feature co-adaptation, where the features of consecutive state-
action tuples learned by the Q-value network become very similar in terms of
their dot product (Section 3). This co-adaptation co-occurs with oscillatory
learning dynamics, and training runs that exhibit feature co-adaptation
typically converge to poorly performing solutions. Even when Q-values are not
overestimated, prolonged training in offline RL can result in performance
degradation as feature co-adaptation increases. To mitigate this co-adaptation
issue, which arises as a result of implicit regularization, we propose an
_explicit regularizer_ that we call DR3 (Section 4). While exactly estimating
and cancelling the effects of the theoretically derived implicit regularizer
is computationally difficult, DR3 provides a simple and tractable
theoretically-inspired approximation that mitigates the issues discussed
above. In practice, DR3 amounts to regularizing the features at consecutive
state-action pairs to be dissimilar in terms of their dot-product similarity.
Empirically, we find that DR3 prevents previously noted pathologies such as
feature rank collapse (Kumar et al., 2021), gives methods that train for
longer and improves performance relative to the base offline RL method
employed in practice.

Our first contribution is the derivation of the implicit regularizer that
arises when training deep net value functions via TD learning, and an
empirical demonstration that it manifests as _feature co-adaptation_ in the
offline deep RL setting. Feature co-adaptation accounts at least in part for
some of the challenges of offline deep RL, including degradation of
performance with prolonged training. Second, we propose a simple and effective
_explicit_ regularizer for offline value-based RL, DR3, which minimizes the
feature similarity between state-action pairs appearing in a bootstrapping
update. DR3 is inspired by the theoretical derivation of the implicit
regularizer, it alleviates co-adaptation and can be easily combined with
modern offline RL methods, such as REM (Agarwal et al., 2020), CQL (Kumar et
al., 2020b), and BRAC (Wu et al., 2019). Empirically, using DR3 in conjunction
with existing offline RL methods provides about 60% performance improvement on
the harder D4RL (Fu et al., 2020) tasks, and 160% and 25% stability gains for
REM and CQL, respectively, on offline RL tasks in 17 Atari 2600 games.
Additionally, we observe large improvements on image-based robotic
manipulation tasks (Singh et al., 2020).

###  2 Preliminaries

The goal in RL is to maximize the long-term discounted reward in an MDP,
defined as $(\mathcal{S},\mathcal{A},R,P,\gamma)$ (Puterman, 1994), with state
space $\mathcal{S}$, action space $\mathcal{A}$, a reward function
$R(\mathbf{s},\mathbf{a})$, dynamics
$P(\mathbf{s}^{\prime}|\mathbf{s},\mathbf{a})$ and a discount factor
$\gamma\in[0,1)$. The Q-function $Q^{\pi}(\mathbf{s},\mathbf{a})$ for a policy
$\pi(\mathbf{a}|\mathbf{s})$ is the expected sum of discounted rewards
obtained by executing action $\mathbf{a}$ at state $\mathbf{s}$ and following
$\pi(\mathbf{a}|\mathbf{s})$ thereafter. $Q^{\pi}(\mathbf{s},\mathbf{a})$ is
the fixed point of
$Q(\mathbf{s},\mathbf{a}):=R(\mathbf{s},\mathbf{a})+\gamma\mathbb{E}_{\mathbf{s}^{\prime}\sim
P(\cdot|\mathbf{s},\mathbf{a}),\mathbf{a}^{\prime}\sim\pi(\cdot|\mathbf{s}^{\prime})}\left[Q(\mathbf{s}^{\prime},\mathbf{a}^{\prime})\right]$.
We study the offline RL setting, where the algorithm must learn a policy only
using a given dataset
$\mathcal{D}=\\{(\mathbf{s}_{i},\mathbf{a}_{i},\mathbf{s}^{\prime}_{i},r_{i})\\}$,
generated from some behavior policy, ${\pi_{\beta}}(\mathbf{a}|\mathbf{s})$,
without active data collection. The Q-function is parameterized with a neural
net with parameters $\theta$. We will denote the penultimate layer of the deep
network (the learned _features_) $\phi_{\theta}(\mathbf{s},\mathbf{a})$, such
that
$Q_{\theta}(\mathbf{s},\mathbf{a})=\mathbf{w}^{T}\phi(\mathbf{s},\mathbf{a})$,
where $\mathbf{w}\in\mathbb{R}^{d}$. Standard deep RL methods (Mnih et al.,
2015; Haarnoja et al., 2018) convert the Bellman equation into a squared
temporal difference (TD) error objective for $Q_{\theta}$:

${\mathcal{L}_{\mathrm{TD}}(\theta)=\sum_{\mathbf{s},\mathbf{a},\mathbf{s}^{\prime}\sim\mathcal{D}}\left(R(\mathbf{s},\mathbf{a})+\gamma\overline{Q}_{\theta}(\mathbf{s}^{\prime},\mathbf{a}^{\prime})-Q_{\theta}(\mathbf{s},\mathbf{a})\right)^{2}},$
(1)

where $\bar{Q}_{\theta}$ is a delayed copy of same Q-network, referred to as
the _target network_ and $\mathbf{a}^{\prime}$ is computed by maximizing the
target Q-function at state $\mathbf{s}^{\prime}$ for Q-learning (i.e., when
computing $Q^{*}$) and by sampling
$\mathbf{a}^{\prime}\sim\pi(\cdot|\mathbf{s})$ when computing the Q-value
$Q^{\pi}$ of a policy $\pi$.

A major problem in offline RL is the issue of distributional shift between the
learned policy and the behavior policy (Levine et al., 2020). Since our goal
is to study the effect of implicit regularization in TD-learning and not
distributional shift, we build on top of existing offline RL methods in our
experiments: CQL (Kumar et al., 2020b), which penalizes erroneous Q-values
during training, REM (Agarwal et al., 2020), which utilizes an ensemble of
Q-functions, and BRAC (Wu et al., 2019), which applies a policy constraint. An
overview of these methods is provided in Appendix E.

###  3 Implicit Regularization in Deep RL via TD-Learning

While the “deadly-triad” (Sutton & Barto, 2018) suggests that training value
function approximators with bootstrapping off-policy can lead to divergence,
modern deep RL algorithms have been able to successfully combine these
properties (van Hasselt et al., 2018). However, making too many TD updates to
the Q-function in offline deep RL is known to sometimes lead to performance
degradation and unlearning, even for otherwise effective modern algorithms (Fu
et al., 2019; Fedus et al., 2020; Agarwal et al., 2020; Kumar et al., 2021).
Such unlearning is not typically observed when training overparameterized
models via supervised learning, so what about TD learning is responsible for
it? We show that one possible explanation behind this pathology is the
implicit regularization induced by minimizing TD error on a deep Q-network.
Our theoretical results suggest that this implicit regularization “co-adapts”
the representations of state-action pairs that appear in a Bellman backup (we
will define this more precisely below). Empirically, this typically manifests
as “co-adapted” features for consecutive state-action tuples, even with
specialized TD-learning algorithms that account for distributional shift, and
this in turn leads to poor final performance both in theory and in practice.
We first provide empirical evidence of this co-adaptation phenomenon in
Section 3.1 (additional evidence in Appendix A.1) and then theoretically
characterize the implicit regularization in TD learning, and discuss how it
can explain the co-adaptation phenomenon in Section 3.2.

Figure 1: Feature dot-products
$\phi(\mathbf{s},\mathbf{a})^{\top}\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})$
increase during training when backing up from _out-of-sample_ but in-
distribution actions (TD-learning: left, Q-learning: right), though the
average Q-value converges and stays relatively constant. Using only seen
state-action pairs for backups (offline SARSA) or not performing Bellman
backups (i.e., supervised regression) avoids this issue, with stable and
relatively low dot products. Left: TD-learning with high feature dot products
eventually destabilizes and produces incorrect Q-values, Right: DQN attains
extremely large feature dot products, despite a relatively stable trend in
Q-values.

####  3.1 Feature Co-Adaptation And How It Relates To Implicit Regularization

In this section, we empirically identify a _feature co-adaptation_ phenomenon
that appears when training value functions via bootstrapping, where the
feature representations of consecutive state-action pairs exhibit a large
value of the dot product
$\phi(\mathbf{s},\mathbf{a})^{\top}\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})$.
Note that feature co-adaptation may arise because of high cosine similarity or
because of high feature norms. Feature co-adaptation appears even when there
is no explicit objective to increase feature similarity.

Experimental setup. We ran supervised regression and three variants of
approximate dynamic programming (ADP) on an offline dataset consisting of 1%
of uniformly-sampled data from the replay buffer of DQN on two Atari games,
previously used in Agarwal et al. (2020). First, for comparison, we trained a
Q-function via supervised regression to Monte-Carlo (MC) return estimates on
the offline dataset to estimate the value of the behavior policy. Then, we
trained variants of ADP which differ in the selection procedure for the action
$\mathbf{a}^{\prime}$ that appears in the target value in
$\mathcal{L}_{\mathrm{TD}}(\theta)$ (Equation 1). The offline SARSA variant
aims to estimate the value of the behavior policy, $Q^{\pi_{\beta}}$, and sets
$\mathbf{a}^{\prime}$ to the actual action observed at the next time step in
the dataset, such that
$(\mathbf{s}^{\prime},\mathbf{a}^{\prime})\in\mathcal{D}$. The TD-learning
variant also aims to estimate the value of the behavior policy, but utilizes
the expectation of the target Q-value over actions $\mathbf{a}^{\prime}$
sampled from the behavior policy $\pi_{\beta}$,
$\mathbf{a}^{\prime}\sim\pi_{\beta}(\cdot|\mathbf{s}^{\prime})$. We do not
have access to the functional form of $\pi_{\beta}$ for the experiment shown
in Figure 1 since the dataset corresponds to the behavior policy induced by
the replay buffer of an online DQN, so we train a model for this policy using
supervised learning. However, we see similar results comparing offline SARSA
and TD-learning on a gridworld domain where we can access the exact functional
form of the behavior policy in Appendix A.6.2. All of the methods so far
estimate $Q^{\pi_{\beta}}$ using different target value estimators. We also
train Q-learning, which chooses the action $\mathbf{a}^{\prime}$ to maximize
the learned Q-function. While Q-learning learns a different Q-function, we can
still compare the relative stability of these methods to gain intuition about
the learning dynamics. In addition to feature dot products
$\phi(\mathbf{s},\mathbf{a})^{\top}\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})$,
we also track the average prediction of the Q-network over the dataset to
measure whether the predictions diverge or are stable in expectation.

Observing feature co-adaptation empirically. As shown in Figure 1 (right), the
average dot product (top row) between features at consecutive state-action
tuples continuously increases for both Q-learning and TD-learning (after
enough gradient steps), whereas it flatlines and converges to a small value
for supervised regression. We might at first think that this is simply a case
of Q-learning failing to converge. However, the bottom row shows that the
average Q-values do in fact converge to a stable value. Despite this, the
optimizer drives the network towards higher feature dot products. There is no
explicit term in the TD error objective that encourages this behavior,
indicating the presence of some implicit regularization phenomenon. This
_implicit_ preference towards maximizing the dot products of features at
consecutive state-action tuples is what we call “feature co-adaptation.”

When does feature co-adaptation emerge? Observe in Figure 1 (right) that the
feature dot products for offline SARSA converge quickly and are relatively
flat, similarly to supervised regression. This indicates that utilizing a
bootstrapped update alone is not responsible for the increasing dot-products
and instability, because while offline SARSA uses backups, it behaves
similarly to supervised MC regression. Unlike offline SARSA, feature co-
adaptation emerges for TD-learning, which is surprising as TD-learning also
aims to estimate the value of the behavior policy, and hence should match
offline SARSA in expectation. The key difference is that while offline SARSA
always utilizes actions $\mathbf{a}^{\prime}$ observed in the training dataset
for the backup, TD-learning may utilize potentially unseen actions
$\mathbf{a}^{\prime}$ in the backup, even though these actions
$\mathbf{a}^{\prime}\sim\pi_{\beta}(\cdot|\mathbf{s}^{\prime})$ are _within_
the distribution of the data-generating policy. This suggests that utilizing
out-of-sample actions in the Bellman backup, even when they are not out-of-
distribution, critically alters the learning dynamics. This is distinct from
the more common observation in offline RL, which attributes training
challenges to out-of-distribution actions (Levine et al., 2020), but not out-
of-sample actions. The theoretical model developed in Section 3.2 will provide
an explanation for this observation with a discussion about how feature co-
adaption caused due to out-of-sample actions can be detrimental in offline RL.

####  3.2 Theoretically Characterizing Implicit Regularization in TD-Learning

Why does feature co-adaptation emerge in TD-learning and what do _out-of-
sample_ actions have to do with it? To answer this question, we theoretically
characterize the implicit regularization effects in TD-learning. We analyze
the learning dynamics of TD learning in the overparameterized regime, where
there are many different parameter vectors $\theta$ that fully minimize the
training set temporal difference error. We base our analysis of TD learning on
the analysis of implicit regularization in supervised learning, previously
developed by Blanc et al. (2020); Damian et al. (2021).

Background. When training an overparameterized $f_{\theta}(\mathbf{x})$ via
supervised regression using the squared loss, denoted by $L$, many different
values of $\theta$ will satisfy $L(\theta)=0$ on the training set due to
overparameterization, but Blanc et al. (2020) show that the dynamics of
stochastic gradient descent will only find fixed points $\theta^{*}$ that
additionally satisfy a condition which can be expressed as
$\nabla_{\theta}R(\theta^{*})=0$, along certain directions (that we will
describe shortly). This function $R(\theta)$ is referred to as the implicit
regularizer. The noisy gradient updates analyzed in this model have the form:

$\theta_{k+1}\leftarrow\theta_{k}-\eta\nabla_{\theta}L(\theta)+\eta\varepsilon_{k},~{}~{}\varepsilon_{k}\sim\mathcal{N}(0,M).$
(2)

Blanc et al. (2020) and Damian et al. (2021) show that some common SGD
techniques fall into this framework, for example, when the regression targets
in supervised learning are corrupted with $\mathcal{N}(0,1)$ label noise, then
the resulting
$M=\sum_{i=1}^{|\mathcal{D}|}\nabla_{\theta}f_{\theta}(\mathbf{x}_{i})\nabla_{\theta}f_{\theta}(\mathbf{x}_{i})^{\top}$
and the induced implicit regularizer $R$ is given by
$R(\theta)=\sum_{i}^{|\mathcal{D}|}||\nabla_{\theta}f_{\theta}(\mathbf{x}_{i})||_{2}^{2}$.
Any solution $\theta^{*}$ found by Equation 2 must satisfy
$\nabla_{\theta}R(\theta^{*})=0$ along directions
$\mathbf{v}\in\mathbb{R}^{|\theta|}$ which lie in the null space of the
Hessian of the loss $\nabla^{2}_{\theta}L(\theta^{*})$ at $\theta^{*}$,
$\mathbf{v}\in\text{Null}(\nabla^{2}_{\theta}L(\theta^{*}))$. The intuition
behind the implicit regularization effect is that along such directions in the
parameter space, the Hessian is unable to contract $\theta_{k}$ when running
noisy gradient updates (Equation 2). Therefore, the only condition that the
noisy gradient updates converge/stabilize at $\theta^{*}$ is given by the
condition that $\nabla R(\theta^{*})=0$. This model corroborates empirical
findings (Mulayoff & Michaeli, 2020; Damian et al., 2021) about the solutions
found by SGD with deep nets, which motivates our use of this framework.

Our setup. Following this framework, we analyze the fixed points of noisy TD-
learning. We consider noisy pseudo-gradient (or semi-gradient) TD updates with
a general noise covariance $M$:

$\displaystyle\theta_{k+1}=\theta_{k}-\eta\underbrace{\left(\sum_{i}\nabla_{\theta}Q(\mathbf{s}_{i},\mathbf{a}_{i})\left(Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})\\!-\\!(r_{i}\\!+\\!\gamma{Q}_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i}))\right)\right)}_{:=g(\theta)}+\eta\varepsilon_{k},~{}~{}\varepsilon_{k}\sim\mathcal{N}(0,M)$
(3)

We use a deterministic policy
$\mathbf{a}^{\prime}_{i}=\pi(\mathbf{s}^{\prime}_{i})$ to simplify exposition.
Following Damian et al. (2021), we can set the noise model $M$ as
$M=\sum_{i}\nabla_{\theta}Q(\mathbf{s}_{i},\mathbf{a}_{i})\nabla_{\theta}Q(\mathbf{s}_{i},\mathbf{a}_{i})^{\top}$,
or utilize a different choice of $M$, but we will derive the general form
first. Let $\theta^{*}$ denote a stationary point of the training TD error,
such that the pseudo-gradient $g(\theta^{*})=0$. Further, we denote the
derivative of $g(\theta)$ w.r.t. $\theta$ as the matrix
$G(\theta)\in\mathbb{R}^{|\theta|\times|\theta|}$, and refer to it as the
_pseudo-Hessian_ : although $G(\theta)$ is not actually the second derivative
of any well-defined objective, since TD updates are not proper gradient
updates, as we will see it will play a similar role to the Hessian in gradient
descent. For brevity, define $G=G(\theta^{*})$, $g=g(\theta^{*})$, $\nabla
G=\nabla_{\theta}G(\theta^{*})\in\mathbb{R}^{|\theta|\times|\theta|\times|\theta|}$,
and let $\lambda_{i}(P)$ denote the $i$-th eigenvalue of matrix $P$, when
arranged in decreasing order of its (complex) magnitude $|\lambda_{i}(P)|$
(note that an eigenvalue can be complex).

Assumptions. To simplify analysis, we assume that matrices $G$ and $M$ (i.e.,
the noise covariance matrix) span the same $n$-dimensional basis in
$d$-dimensional space, where $d$ is the number of parameters and $n$ is the
number of datapoints, and $n\ll d$ due to overparameterization. We also
require $\theta^{*}$ to satisfy a technical criterion that requires
approximate alignment between the eigenspaces of $G$ and the gradient of the
Q-function, without which noisy TD may not be stable at $\theta^{*}$. We
summarize all the assumptions in Appendix C, and present the resulting
regularizer below.

######  Theorem 3.1 (Implicit regularizer at TD fixed points).

Under the assumptions so far, a fixed point of TD-learning, $\theta^{*}$,
where $Q_{\theta^{*}}(\mathbf{s}_{i},\mathbf{a}_{i})=r_{i}+\gamma
Q_{\theta^{*}}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$ for every
$(\mathbf{s}_{i},\mathbf{a}_{i},\mathbf{s}^{\prime}_{i})\in\mathcal{D}$ is
stable (atttractive) if: (1) it satisfies $\mathrm{Re}(\lambda_{i}(G))\geq
0,\forall i$ and $\mathrm{Re}(\lambda_{i}(G))>0$ if
$|\mathrm{Imag}(\lambda_{i}(G))|>0$, and (2) along directions
$\mathbf{v}\in\mathbb{R}^{\text{dim}(\theta)},\mathbf{v}\in\text{Null}(G)$,
$\theta^{*}$ is the stationary point of the induced implicit regularizer:

$\displaystyle R_{\mathrm{TD}}(\theta)$
$\displaystyle=\underbracket{\eta\sum_{i=1}^{|\mathcal{D}|}\nabla
Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})^{\top}\Sigma^{*}_{M}\nabla
Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})}_{\text{the implicit regularizer for
noisy GD in supervised learning}}$
$\displaystyle~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}~{}-\underbracket{\eta\gamma\sum_{i=1}^{|\mathcal{D}|}\mathrm{trace}\left(\left[\left[\nabla
Q_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})^{\top}\right]\right]^{\top}\Sigma_{M}^{*}\nabla
Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})\right)}_{\text{additional term in TD
learning}},$ (4)

where $(\mathbf{s}_{i},\mathbf{a}_{i})$ and
$(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$ denote state-action pairs
that appear together in a Bellman update, $[[\square]]$ denotes the stop-
gradient function, which does not pass partial gradients w.r.t. $\theta$ into
$\square$. $\Sigma^{*}_{M}$ is the fixed point of the discrete Lyapunov
equation: $\Sigma^{*}_{M}:=(I-\eta G)\Sigma^{*}_{M}(I-\eta
G)^{\top}+\eta^{2}M$.

A proof of Theorem 3.1 is provided in Appendix C. Next, we explain the
intuition behind this result and provide a proof sketch. To derive the induced
implicit regularizer for a stable fixed point $\theta^{*}$ of TD error, we
study the learning dynamics of noisy TD learning (Equation 3) initialized at
$\theta^{*}$, and derive conditions under which this noisy update would stay
close to $\theta^{*}$ with multiple updates. This gives rise to the two
conditions shown in Theorem 3.1 which can be understood as controlling
stability in mutually exclusive directions in the parameter space. If
condition (1) is not satisfied, then even under-parameterized TD will diverge
away from $\theta^{*}$, since $I-\eta G$ would be a non-contraction as the
spectral radius, $\rho(I-\eta G)\geq 1$ in that case. Thus,
$\theta_{k}-\theta^{*}$ will grow or not decrease in some direction. When (1)
is satisfied for all directions in the parameter space, there are still
directions where both the real and imaginary parts of the eigenvalue
$\lambda_{i}(G)$ are $0$ due to overparameterization111To see why this is the
case, note that $\text{rank}(G)\leq|\mathcal{D}|\ll\text{dim}(\theta)$, and so
some eigenvalues of $G$ are $0$.. In such directions, learning is governed by
the projection of the noise under the tensor $\nabla G$, which appears in the
Taylor expansion of $\theta_{k}-\theta^{*}$ around the point $\theta^{*}$:

$\displaystyle\theta_{k+1}=\theta_{k}-\eta\left(g+G(\theta_{k}-\theta^{*})+\frac{1}{2}\nabla
G[\theta_{k}-\theta^{*},\theta_{k}-\theta^{*}]\right)+\varepsilon_{k},~{}~{}\varepsilon_{k}\sim\mathcal{N}(0,M)$
(5) $\displaystyle\implies$ $\displaystyle\nu_{k+1}=(I-\eta
G)\nu_{k}-\frac{\eta}{2}\nabla G[\nu_{k},\nu_{k}]+\varepsilon_{k},$ (6)

where we reparameterize in terms of $\nu_{k}:=\theta_{k}-\theta^{*}$. The
proof shows that $\theta^{*}$ is stable if it is a stationary point of the
implicit regularizer $R_{\mathrm{TD}}$ (condition (2)), which ensures that
total noise (i.e., accumulated $\varepsilon_{k}$ over iterations $k$)
accumulated by $\nabla G$ does not lead to a large deviation in $\nu_{k}$ in
directions where $I-\eta G$ does not contract.

Interpretation of Theorem 3.1. While the choice of the noise model $M$ will
change the form of the implicit regularizer, in practice, the form of $M$ is
not known as this corresponds to the noise induced via SGD. We can consider
choices of $M$ for interpretation, but Theorem 3.1 is easy to qualitatively
interpret for $M$ such that $\Sigma^{*}_{M}=I$. In this case, we find that the
implicit preference towards local minima of $R_{\mathrm{TD}}(\theta)$ can
explain feature co-adaptation. In this case, the regularizer takes a simpler
form:

$\displaystyle R_{\mathrm{TD}}(\theta):=\sum_{i}||\nabla
Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})||_{2}^{2}-\gamma\nabla
Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})\nabla[[Q_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})]].$

The first term is equal to the squared per-datapoint gradient norm, which is
same as the implicit regularizer in supervised learning obtained by Blanc et
al. (2020); Damian et al. (2021) with label noise. However,
$R_{\mathrm{TD}}(\theta)$ additionally includes a second term that is equal to
the dot product of the gradient of the Q-function at the current and next
states,
$\nabla_{\theta}Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})^{\top}\nabla_{\theta}Q_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$,
and thus this term is effectively _maximized_. When restricted to the last-
layer parameters of a neural network, this term is equal to the dot product of
the features at consecutive state-action tuples:
$\sum_{i}\nabla_{\theta}Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})^{\top}\nabla_{\theta}Q_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})=\sum_{i}\phi(\mathbf{s}_{i},\mathbf{a}_{i})^{\top}\phi(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$.
The tendency to maximize this quantity to attain a local minimizer of the
implicit regularizer corroborates the empirical findings of increased dot
product in Section 3.1.

Explaining the difference between utilizing seen and unseen actions in the
backup. If all state-action pairs
$(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$ appearing on the right-
hand-side of the Bellman update also appear in the dataset $\mathcal{D}$, as
in the case of offline SARSA (Figure 1), the preference to increase dot
products will be balanced by the affinity to reduce gradient norm (first term
of $R_{\mathrm{TD}}(\theta)$ when $\Sigma^{*}_{M}=I$): for example, for
offline SARSA, when $(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$ are
permutations of $(\mathbf{s}_{i},\mathbf{a}_{i})$, $R_{\mathrm{TD}}$ is lower
bounded by
$(1-\gamma)\sum_{i}||\nabla_{\theta}Q_{\theta}(\mathbf{x}_{i})||_{2}^{2}$ and
hence minimizing $R_{\mathrm{TD}}(\theta)$ would minimize the feature norm
instead of maximizing dot products. This also corresponds to the implicit
regularizer we would obtain when training Q-functions via supervised learning
and hence, our analysis predicts that offline SARSA with in-sample actions
(i.e., when $(\mathbf{s}^{\prime},\mathbf{a}^{\prime})\in\mathcal{D}$) would
behave similarly to supervised regression.

However, the regularizer behaves very differently when unseen state-action
pairs $(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$ appear only on the
right-hand-side of the backup. This happens with any algorithm where
$\mathbf{a}^{\prime}$ is not the dataset action, which is the case for all
deep RL algorithms that compute target values by selecting
$\mathbf{a}^{\prime}$ according to the current policy. In this case, we expect
the dot product of gradients at $(\mathbf{s},\mathbf{a})$ and
$(\mathbf{s}^{\prime},\mathbf{a}^{\prime})$ to be large at any attractive
fixed point, since this minimizes $R_{\mathrm{TD}}(\theta)$. This is precisely
a form of co-adaptation: gradients at out-of-sample state-action tuples are
highly similar to gradients at observed state-action pairs measured by the dot
product. This observation is also supported by the analysis in Section 3.1.
Finally, note that the choice of $M$ is a modelling assumption, and to derive
our explicit regularizer, later in the paper, we will make a simplifying
choice of $M$. However, we also empirically verify that a different choice of
$M$, given by label noise, works well.

Figure 2: Even when current offline RL algorithms are initialized at a high-
performing checkpoint that attains small feature dot products, feature dot
products increase with further training and the performance degrades.

Why is implicit regularization detrimental to policy performance? To answer
this question, we present theoretical and empirical evidence that illustrates
the adverse effects of this implicit regularizer. Empirically, we ran two
algorithms, DQN and CQL, initialized from a high-performing Q-function
checkpoint, which attains relatively small feature dot products (i.e., the
second term of $R_{\mathrm{TD}}(\theta)$ is small). Our goal is to see if TD
updates starting from such a “good” initialization still stay around it or
diverge to poorer solutions. Our theoretical analysis in Section 3.2 would
predict that TD learning would destabilize from such a solution, since it
would not be a stable fixed point. Indeed, as shown in Figure 2, the policy
immediately degrades, and the the dot-product similarities start to increase.
This even happens with CQL, which explicitly corrects for distributional shift
confounds, implying that the performance drop cannot be directly explained by
the typical out-of-distribution action explanations. To investigate the
reasons behind this drop, we also measured the training loss function values
for these algorithms (i.e., TD error for DQN and TD error + CQL regularizer
for CQL) and find in Figure 2 that the loss values are generally small for
both CQL and DQN. This indicates that the preference to increase dot products
is not explained by an inability to minimize TD error. In Appendix A.7, we
show that this drop in performance when starting from good solutions can be
effectively mitigated with our proposed DR3 explicit regularizer for both DQN
and CQL. Thus we find that not only standard TD learning degrades from a good
solution in favor of increasing feature dot products, but keeping small dot
products enables these algorithms to remain stable near the good solution.

To motivate why co-adapted features can lead to poor performance in TD-
learning, we study the convergence of linear TD-learning on co-adapted
features. Our theoretical result characterizes a lower bound on the feature
dot products in terms of the feature norms for state-action pairs in the
dataset $\mathcal{D}$, which if satisfied, will inhibit convergence:

######  Proposition 3.2 (TD-learning on co-adapted features).

Assume that the features
$\Phi=[\phi(\mathbf{s},\mathbf{a})]_{\mathbf{s},\mathbf{a}}$ are used for
linear TD-learning. Then, if
$\sum_{\mathbf{s},\mathbf{a},\mathbf{s}^{\prime}\in\mathcal{D}}\phi(\mathbf{s},\mathbf{a})^{\top}\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})\geq\frac{1}{\gamma}\sum_{\mathbf{s},\mathbf{a}\in\mathcal{D}}\phi(\mathbf{s},\mathbf{a})^{\top}\phi(\mathbf{s},\mathbf{a})$,
linear TD-learning using features $\Phi$ will not converge.

A proof of Proposition 3.2 is provided in Appendix D and it relies on a
stability analysis of linear TD. While features change during training for TD-
learning with neural networks, and arguably linear TD is a simple model to
study consequences of co-adapted features, even in this simple linear setting,
Proposition 3.2 indicates that TD-learning may be non-convergent as a result
of co-adaptation.

Comparison to implicit regularization in TD learning with linear function
approximation. Running stochastic gradient descent in overparameterized linear
regression finds solutions with the smallest $\ell_{2}$ norm, which is often
regarded as the implicit regularizer. Based on this observation, one might
wonder how our derived implicit regularizer relates to minimum norm solutions
attained by gradient descent in overparameterized linear TD learning. The
implicit regularizer we obtain in Equation 3.1 would be a constant,
independent of the parameter vector $\theta$ for linear TD learning. Thus our
regularization specifically captures the effect of SGD on non-linear function
approximators, which are absent when studying linear function approximation.

Takeaways. We summarize the key takeaways from our theoretical analysis below:

  * •

_The implicit regularizer at TD fixed points_ is shown in Equation 3.1. The
first term corresponds to the regularizer for SGD in supervised learning,
while the second term that is unique to TD and leads to an (undesirable)
increase in gradient or feature dot products.

  * •

_Out-of-sample actions exacerbate the implicit regularization effect,_ since
feature dot products can be easily increased when out-of-sample actions, which
do not appear in the dataset, are used to compute Bellman targets.

  * •

The implicit regularizer in Equation 3.1 is induced via a mechanism unique to
non-linear Q-functions, different from overparameterized, linear TD-learning.

###  4 DR3: Explicit Regularization for Deep TD-Learning

Since the implicit regularization effects in TD-learning can lead to feature
co-adaptation, which in turn is correlated with poor performance, can we
instead derive an _explicit_ regularizer to alleviate this issue? Inspired by
the analysis in the previous section, we will propose an _explicit_
regularizer that attempts to counteract the second term in Equation 3.1, which
would otherwise lead to co-adaptation and poor representations. The _explicit_
regularizer that offsets the difference between the two implicit regularizers
is given by:
$\Delta(\theta)=\sum_{i}\mathrm{trace}\left[\Sigma^{*\top}_{M}\nabla_{\theta}Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})\nabla_{\theta}Q_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})^{\top}\right]$,
which represents the second term of $R_{\mathrm{TD}}(\theta)$. Note that we
drop the stop gradient on
$Q_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$ in
$\Delta(\theta)$, as it performs slightly better in practice (Table A.1),
although as shown in that Table, the version with the stop gradient also
significantly improves over the base method. The first term of
$R_{\mathrm{TD}}(\theta)$ corresponds to the regularizer from supervised
learning. Our proposed method, DR3, simply combines approximations to
$\Delta(\theta)$ with various offline RL algorithms. For any offline RL
algorithm, Alg, with objective $\mathcal{L}_{\textsc{Alg}}(\theta)$, the
training objective with DR3 is given by:
$\mathcal{L}(\theta):=\mathcal{L}_{\textsc{Alg}}(\theta)+c_{0}\Delta(\theta)$,
where $c_{0}$ is the DR3 coefficient. See Appendix E.3 for details on how we
tune $c_{0}$ in this paper.

Practical version of DR3. In order to practically instantiate DR3, we need to
choose a particular noise model $M$. In general, it is not possible to know
beforehand the “correct” choice of $M$ (Equation 3), even in supervised
learning, as this is a complicated function of the data distribution, neural
network architecture and initialization. Therefore, we instantiate DR3 with
two heuristic choices of $M$: (i) $M$ induced by label noise studied in prior
work for supervised learning and for which we need to run a computationally
heavy fixed-point computation for $M$, and (ii) a simpler alternative that
sets $\Sigma^{*}_{M}=I$. We find that both of these variants generally perform
well empirically (Figure 6), and improve over the base offline RL method, and
so we utilize (ii) in practice due to low computational costs. Additionally,
because computing and backpropagating through per-example gradient dot
products is slow, we instead approximate $\Delta(\theta)$ with the
contribution only from the last layer parameters (_i.e._ ,
$\sum_{i}\nabla_{\mathbf{w}}Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})^{\top}\nabla_{\mathbf{w}}Q_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$),
similarly to tractable Bayesian neural nets. As shown in Appendix A.5, the
practical version of DR3 performs similarly to the label-noise version.

$\text{{Explicit DR3
regularizer}}:~{}~{}~{}~{}~{}~{}~{}~{}\overline{\mathcal{R}}_{\mathrm{exp}}(\theta)=\sum_{i\in\mathcal{D}}\phi(\mathbf{s}_{i},\mathbf{a}_{i})^{\top}\phi(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i}).$
(7)

###  5 Related Work

Prior analyses of the learning dynamics in RL has focused primarily on
analyzing error propagation in tabular or linear settings (_e.g._ , Chen &
Jiang, 2019; Duan et al., 2020; Xie & Jiang, 2020; Wang et al., 2021a; b;
Farahmand et al., 2010; De Farias, 2002), understanding instabilities in deep
RL (Achiam et al., 2019; Bengio et al., 2020; Kumar et al., 2020a; Van Hasselt
et al., 2018) and deriving weighted TD updates that enjoy convergence
guarantees (Maei et al., 2009; Mahmood et al., 2015; Sutton et al., 2016), but
these methods do not reason about implicit regularization or any form of
representation learning. Ghosh & Bellemare (2020) focuses on understanding the
stability of TD-learning in underparameterized linear settings, whereas our
focus is on the overparameterized setting, when optimizing TD error and
learning representations via SGD. Kumar et al. (2021) studies the learning
dynamics of Q-learning and observes that the rank of the feature matrix,
$\Phi$, drops during training. While this observation is related, our analysis
characterizes the implicit preference of learning towards feature co-
adaptation (Theorem 3.1) on out-of-sample actions as the primary culprit for
aliasing. Additionally, while the goal of our work is not to increase
$\text{srank}(\Phi)$, utilizing DR3 not only outperforms the
$\text{srank}(\Phi)$ penalty in Kumar et al. (2021) by more than 100%, but it
also alleviates rank collapse, with no apparent term that explicitly enforces
high rank values. Somewhat related to DR3, Durugkar & Stone (2018); Pohlen et
al. (2018) heuristically constrain gradients of TD to prevent changes in
target Q-values to prevent divergence. Contrary to such heuristic approaches,
DR3 is inspired from a theoretical model of implicit regularization, and does
not prevent changes in target values, but rather reduces feature dot products.

###  6 Experimental Evaluation of DR3

Our experiments aim to evaluate the extent to which DR3 improves performance
in offline RL in practice, and to study its effect on prior observations of
rank collapse. To this end, we investigate if DR3 improves offline RL
performance and stability on three offline RL benchmarks: Atari 2600 games
with discrete actions (Agarwal et al., 2020), continuous control tasks from
D4RL (Fu et al., 2020), and image-based robotic manipulation tasks (Singh et
al., 2020).

Following prior work (Fu et al., 2020; Gulcehre et al., 2020), we evaluate DR3
in terms of final offline RL performance after a given number of iterations.
Additionally, we report _training stability_ , which is important in practice
as offline RL does not admit cheap validation of trained policies for model
selection. To evaluate stability, we train for a large number of gradient
steps (2-3x longer than prior work) and either report the average performance
over the course of training or the final performance at the end of training.
We expect that a stable method that does not unlearn with more gradient steps,
should have better average performance, as compared to a method that attains
good peak performance but degrades with more training. See Appendix E for
further details.

Offline RL on Atari 2600 games. We compare DR3 to prior offline RL methods on
a set of offline Atari datasets of varying sizes and quality, akin to Agarwal
et al. (2020); Kumar et al. (2021). We evaluated on three datasets: (1) 1% and
5% samples drawn uniformly at random from DQN replay; (2) a dataset with more
suboptimal data consisting of the first 10% samples observed by an online DQN.
Following Agarwal et al. (2021), we report the interquartile mean (IQM)
normalized scores across 17 games over the course of training in Figure 4 and
report the IQM average performance in Table 1. Observe that combining DR3 with
modern offline RL methods (CQL, REM) attains the best final and average
performance across the 17 Atari games tested on, directly improving upon prior
methods across all the datasets. When DR3 is used in conjunction with REM, it
prevents severe unlearning and performance degradation with more training. CQL
+ DR3 improves by 20% over CQL on final performance and attains 25% better
average performance. While DR3 is not unequivocally “stable”, as its
performance also degrades relative to the peak it achieves (Figure 4), it is
more stable relative to base offline RL algorithms. We also compare DR3 to the
$\mathrm{srank}(\Phi)$ penalty proposed to counter rank collapse (Kumar et
al., 2021). Directly taking median normalized score improvements reported by
Kumar et al. (2021), CQL + DR3 improves by over 2x (31.5%) over naïve CQL
relative to the srank penalty (14.1%), indicating DR3’s efficacy.

Offline RL on robotic manipulation from images. Next, we aim to evaluate the
efficacy of DR3 on two image-based robotic manipulation tasks (Singh et al.,
2020) (visualized on the right) that require composition of skills (e.g.,
opening a drawer, closing a drawer, picking an obstructive object, placing an
object, etc.) over extended horizons using only a sparse 0-1 reward. As shown
in Figure 4, combining DR3 with COG not only improves over COG, but also
learns faster and attains a better average performance.

Table 1: IQM normalized average performance (training stability) across 17 games, with 95% CIs in parenthesis, after 6.5M gradient steps for the 1% setting and 12.5M gradient steps for the 5%, 10% settings. Individual performances reported in Tables F.4-F.9. DR3 improves the stability over both CQL and REM.  Data | CQL | CQL + DR3 | REM | REM + DR3
---|---|---|---|---
1% | 43.7 (39.6, 48.6) |  56.9 (52.5, 61.2) | 4.0 (3.3, 4.8) |  16.5 (14.5, 18.6)
5% | 78.1 (74.5, 82.4) |  105.7 (101.9, 110.9) | 25.9 (23.4, 28.8) |  60.2 (55.8, 65.1)
10% | 59.3 (56.4, 61.9) |  65.8 (63.3, 68.3) | 53.3 (51.4, 55.3) |  73.8 (69.3, 78)

Figure 3: Performance of DR3 + COG on two manipulation tasks using only 5% and
25% of the data used by Singh et al. (2020) to make these more challenging.
COG + DR3 outperforms COG in training and attains higher average and final
performance.

Figure 4: Normalized performance across 17 Atari games for REM + DR3 (top),
CQL + DR3 (bottom). x-axis represents _gradient steps_ ; no new data is
collected. While naïve REM suffers from a degradation in performance with more
training, REM + DR3 not only remains generally stable with more training, but
also attains higher final performance. CQL + DR3 attains higher performance
than CQL. We report IQM with 95% stratified bootstrap CIs (Agarwal et al.,
2021).

Offline RL on D4RL tasks. Finally, we evaluate DR3 in conjunction with CQL on
the antmaze-v2 domain in D4RL (Fu et al., 2020). To assess if DR3 is stable
and able to prevent unlearning that eventually appears in CQL, we trained
CQL+DR3 for 9x longer: 2M and 3M steps with 3x higher learning rate. This is
different from prior works (Fu et al., 2020) that report performance at the
end of 1M steps. Observe in Table 2, that CQL + DR3 outperforms CQL
(statistical significance shown in Appendix A.8), indicating that DR3
significantly improves CQL. We also evalute DR3 on kitchen domains in D4RL in
Appendix A.8, where we also find that DR3 improves CQL.

Table 2: Performance of CQL, CQL + DR3 after 2M and 3M gradient steps with a
learning rate of 3e-4 for the Q-function averaged over 3 seeds. This is
training for 6x and 9x longer compared to CQL defaults. Observe that CQL + DR3
outperforms CQL at 2M and 3M steps, indicating is efficacy in preventing
unlearning. We present the statistical significance of these results in
Appendix A.8.
D4RL Task |  CQL (2M) |  CQL + DR3 (2M) |  CQL (3M) |  CQL + DR3 (3M)
---|---|---|---|---
antmaze-umaze-v2 | 84.00 $\pm$ 2.67 | 85.33 $\pm$ 4.16 | 87.00 $\pm$ 1.73 | 90.00 $\pm$ 4.00
antmaze-umaze-diverse-v2 | 45.67 $\pm$ 8.50 | 40.67 $\pm$ 11.84 | 36.33 $\pm$ 7.09 | 52.00 $\pm$ 11.26
antmaze-medium-play-v2 | 24.00 $\pm$ 28.16 | 73.00 $\pm$ 4.00 | 16.00 $\pm$ 26.85 | 71.33 $\pm$ 1.52
antmaze-medium-diverse-v2 | 32.67 $\pm$ 9.29 | 67.00 $\pm$ 2.00 | 48.33 $\pm$ 6.11 | 61.67 $\pm$ 3.21
antmaze-large-play-v2 | 3.33 $\pm$ 2.51 | 28.00 $\pm$ 4.35 | 0.33 $\pm$ 0.57 | 26.33 $\pm$ 11.93
antmaze-large-diverse-v2 | 1.33 $\pm$ 2.30 | 25.67 $\pm$ 0.57 | 0.00 $\pm$ 0.00 | 28.33 $\pm$ 1.52

Finally, we also compare CQL+DR3 and CQL in terms of performance and stability
on MuJoCo tasks previously studied in Kumar et al. (2021) in Appendix A.3.
These tasks are constructed by uniformly subsampling transitions from the
full-replay-v2 MuJoCo datasets in D4RL and are much harder than the typical
Gym-MuJoCo tasks from Fu et al. (2020) because succeeding on these tasks
critically relies on estimating accurate Q-values for out-of-sample actions
and all actions at certain states are out-of-sample. As shown in Appendix A.3,
CQL+DR3 is significantly more stable, and does not unlearn with more training,
unlike CQL whose performance degrades very quickly. We also evaluate DR3 in
conjunction with BRAC (Wu et al., 2019), a policy constraint method, and find
that BRAC+DR3 improves over BRAC in 13.8 median normalized performance (Table
F.2). To summarize, these results indicate that DR3 is a versatile explicit
regularizer that improves performance and stability of a wide range of offline
RL methods, including conservative methods (e.g, CQL, COG), policy constraint
methods (e.g., BRAC) and ensemble-based methods (e.g., REM).

Figure 5: Trend of effective rank, $\mathrm{srank}(\Phi)$ of features $\Phi$
learned by the Q-function when trained with TD error (red, “Without DR3”) and
with TD error + DR3 (blue, “With DR3”) on three Atari games using the 5%
dataset. Note that DR3 alleviates rank collapse, without explicitly aiming to.

DR3 does not suffer from rank collapse. Prior work (Kumar et al., 2021) has
shown that implicit regularization can lead to a rank collapse issue in TD-
learning, preventing Q-networks from using full capacity. To see if DR3
addresses the rank collapse issue, we follow Kumar et al. (2021) and plot the
effective rank of learned features with DR3 in Figure 5 (DQN, REM in Appendix
A.4). While the value of the effective rank decreases during training with
naïve bootstrapping, we find that rank of DR3 features typically does not
collapse, despite no explicit term encouraging this. Finally, we test the
robustness/sensitivity of each layer in the learned Q-network to re-
initialization (Zhang et al., 2019) during training and find that DR3 alters
the features to behave similarly to supervised learning (Figure A.2).

Figure 6:  Comparing DR3 regularizers for our simplifying choice of $M$ and
$M$ induced by label noise, with base CQL and DQN algorithms. Note that both
of these penalties when applied over CQL improve performance.

Comparing explicit regularizers for different choices of noise covariance $M$.
Finally, we investigate the behavior of different implicit regularizers
derived via two choices of $M$ in Equation 3.1 and the corresponding explicit
regularizers. While the explicit regularizer we use in practice is a
simplifying choice that works well, another choice of $M$ is the covariance
matrix induced by label noise, which requires explicit computation of
$\Sigma_{M}^{*}$. Observe in Figure 6 that the explicit regularizer for our
simplifying choice is not worse than the different choice of $M$. This
justifies utilizing our simplified, heuristic choice of setting
$\Sigma_{M}^{*}=I$ in practice. Results on five Atari games are shown in
Appendix A.5.

###  7 Discussion

We characterized the implicit preference of TD-learning towards solutions that
maximally co-adapt gradients (or features) at consecutive state-action tuples
that appear in Bellman backup. This regularization effect is exacerbated when
out-of-sample state-action samples are used for the Bellman backup and it can
lead to poor policy performance. Inspired by the theory, we propose a
practical explicit regularizer, DR3 that aims to counteracts this implicit
regularizer. DR3 yields substantial improvements in stability and performance
on a wide range of offline RL problems. We believe that understanding the
learning dynamics of deep Q-learning and the induced implicit regularization
will lead to more robust and stable deep RL algorithms. Furthermore, this
understanding can help us to predict the instability issues in value-based RL
methods in advance, which can inspire cross-validation and model selection
strategies, an important, open challenge in offline RL, for which existing
off-policy evaluation techniques are not practically sufficient (Fu et al.,
2021). We also note that our analysis does not consider the online RL setting
with non-stationary data distributions, and extending our theory and DR3 to
online RL is an interesting avenue for future work.

### Acknowledgements

We thank Dibya Ghosh, Xinyang Geng, Dale Schuurmans, Marc Bellemare, Pablo
Castro and Ofir Nachum for informative discussions, discussions on
experimental setup and for providing feedback on an early version of this
paper. We thank the members of RAIL at UC Berkeley for their support and
suggestions. We thank anonymous reviewers for feedback on an early version of
this paper. This research is funded in part by the DARPA Assured Autonomy
Program and in part, by compute resources from Microsoft Azure and Google
Cloud. TM acknowledges support of Google Faculty Award, NSF IIS 2045685, the
Sloan fellowship, and JD.com.

### References

  * Achiam et al. (2019) Joshua Achiam, Ethan Knight, and Pieter Abbeel.  Towards characterizing divergence in deep q-learning.  _arXiv preprint arXiv:1903.08894_ , 2019.
  * Agarwal et al. (2020) Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi.  An optimistic perspective on offline reinforcement learning.  In _International Conference on Machine Learning (ICML)_ , 2020.
  * Agarwal et al. (2021) Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc G Bellemare.  Deep reinforcement learning at the edge of the statistical precipice.  _Advances in Neural Information Processing Systems_ , 2021.
  * Arora et al. (2018) Sanjeev Arora, Nadav Cohen, and Elad Hazan.  On the optimization of deep networks: Implicit acceleration by overparameterization.  _arXiv preprint arXiv:1802.06509_ , 2018.
  * Bengio et al. (2020) Emmanuel Bengio, Joelle Pineau, and Doina Precup.  Interference and generalization in temporal difference learning.  _arXiv preprint arXiv:2003.06350_ , 2020.
  * Blanc et al. (2020) Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant.  Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process.  In _Conference on learning theory_ , pp. 483–513. PMLR, 2020.
  * Bouthillier et al. (2021) Xavier Bouthillier, Pierre Delaunay, Mirko Bronzi, Assya Trofimov, Brennan Nichyporuk, Justin Szeto, Nazanin Mohammadi Sepahvand, Edward Raff, Kanika Madan, Vikram Voleti, et al.  Accounting for variance in machine learning benchmarks.  _Proceedings of Machine Learning and Systems_ , 3, 2021.
  * Chatterji et al. (2019) Niladri S Chatterji, Behnam Neyshabur, and Hanie Sedghi.  The intriguing role of module criticality in the generalization of deep networks.  _arXiv preprint arXiv:1912.00528_ , 2019.
  * Chen & Jiang (2019) Jinglin Chen and Nan Jiang.  Information-theoretic considerations in batch reinforcement learning.  _ICML_ , 2019.
  * Chen & He (2020) Xinlei Chen and Kaiming He.  Exploring simple siamese representation learning.  _arXiv preprint arXiv:2011.10566_ , 2020.
  * Damian et al. (2021) Alex Damian, Tengyu Ma, and Jason Lee.  Label noise sgd provably prefers flat global minimizers.  _arXiv preprint arXiv:2106.06530_ , 2021.
  * De Farias (2002) Daniela Pucci De Farias.  _The linear programming approach to approximate dynamic programming: Theory and application_.  PhD thesis, 2002.
  * Duan et al. (2020) Yaqi Duan, Zeyu Jia, and Mengdi Wang.  Minimax-optimal off-policy evaluation with linear function approximation.  In _International Conference on Machine Learning_ , pp. 2701–2709. PMLR, 2020.
  * Durugkar & Stone (2018) Ishan Durugkar and Peter Stone.  Td learning with constrained gradients.  2018\.
  * Farahmand et al. (2010) Amir-massoud Farahmand, Csaba Szepesvári, and Rémi Munos.  Error propagation for approximate policy and value iteration.  In _Advances in Neural Information Processing Systems (NIPS)_ , 2010\.
  * Fedus et al. (2020) William Fedus, Prajit Ramachandran, Rishabh Agarwal, Yoshua Bengio, Hugo Larochelle, Mark Rowland, and Will Dabney.  Revisiting fundamentals of experience replay.  _arXiv preprint arXiv:2007.06700_ , 2020.
  * Fu et al. (2019) Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine.  Diagnosing bottlenecks in deep Q-learning algorithms.  _arXiv preprint arXiv:1902.10250_ , 2019.
  * Fu et al. (2020) Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_ , 2020.
  * Fu et al. (2021) Justin Fu, Mohammad Norouzi, Ofir Nachum, George Tucker, ziyu wang, Alexander Novikov, Mengjiao Yang, Michael R Zhang, Yutian Chen, Aviral Kumar, Cosmin Paduraru, Sergey Levine, and Thomas Paine.  Benchmarks for deep off-policy evaluation.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=kWSeGEeHvF8.
  * Ghosh & Bellemare (2020) Dibya Ghosh and Marc G Bellemare.  Representations for stable off-policy reinforcement learning.  _arXiv preprint arXiv:2007.05520_ , 2020.
  * Grill et al. (2020) Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al.  Bootstrap your own latent: A new approach to self-supervised learning.  _arXiv preprint arXiv:2006.07733_ , 2020.
  * Gulcehre et al. (2020) Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Tom Le Paine, Sergio Gómez Colmenarejo, Konrad Zolna, Rishabh Agarwal, Josh Merel, Daniel Mankowitz, Cosmin Paduraru, et al.  Rl unplugged: Benchmarks for offline reinforcement learning.  2020\.
  * Gunasekar et al. (2017) Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro.  Implicit regularization in matrix factorization.  In _Advances in Neural Information Processing Systems_ , pp. 6151–6159, 2017.
  * Haarnoja et al. (2018) Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  _CoRR_ , abs/1801.01290, 2018.  URL http://arxiv.org/abs/1801.01290.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clement Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  In _Advances in Neural Information Processing Systems 31_. 2018.
  * Kumar et al. (2020a) Aviral Kumar, Abhishek Gupta, and Sergey Levine.  Discor: Corrective feedback in reinforcement learning via distribution correction.  _arXiv preprint arXiv:2003.07305_ , 2020a.
  * Kumar et al. (2020b) Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  _arXiv preprint arXiv:2006.04779_ , 2020b.
  * Kumar et al. (2021) Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine.  Implicit under-parameterization inhibits data-efficient deep reinforcement learning.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=O9bnihsFfXU.
  * Levine et al. (2020) Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  _arXiv preprint arXiv:2005.01643_ , 2020.
  * Li et al. (2019) Yuanzhi Li, Colin Wei, and Tengyu Ma.  Towards explaining the regularization effect of initial large learning rate in training neural networks.  In _Advances in Neural Information Processing Systems_ , pp. 11674–11685, 2019.
  * Maei et al. (2009) Hamid R. Maei, Csaba Szepesvári, Shalabh Bhatnagar, Doina Precup, David Silver, and Richard S. Sutton.  Convergent temporal-difference learning with arbitrary smooth function approximation.  In _Proceedings of the 22nd International Conference on Neural Information Processing Systems_ , 2009.
  * Mahmood et al. (2015) A Rupam Mahmood, Huizhen Yu, Martha White, and Richard S Sutton.  Emphatic temporal-difference learning.  _arXiv preprint arXiv:1507.01569_ , 2015.
  * Mnih et al. (2015) Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis.  Human-level control through deep reinforcement learning.  _Nature_ , 518(7540):529–533, feb 2015.  ISSN 0028-0836.
  * Mulayoff & Michaeli (2020) Rotem Mulayoff and Tomer Michaeli.  Unique properties of flat minima in deep networks.  In _International Conference on Machine Learning_ , pp. 7108–7118. PMLR, 2020.
  * Pohlen et al. (2018) Tobias Pohlen, Bilal Piot, Todd Hester, Mohammad Gheshlaghi Azar, Dan Horgan, David Budden, Gabriel Barth-Maron, Hado Van Hasselt, John Quan, Mel Večerík, et al.  Observe and look further: Achieving consistent performance on atari.  _arXiv preprint arXiv:1805.11593_ , 2018.
  * Puterman (1994) Martin L Puterman.  _Markov Decision Processes: Discrete Stochastic Dynamic Programming_.  John Wiley & Sons, Inc., 1994.
  * Riedmiller (2005) Martin Riedmiller.  Neural fitted q iteration–first experiences with a data efficient neural reinforcement learning method.  In _European Conference on Machine Learning_ , pp. 317–328. Springer, 2005.
  * Singh et al. (2020) Avi Singh, Albert Yu, Jonathan Yang, Jesse Zhang, Aviral Kumar, and Sergey Levine.  Cog: Connecting new skills to past experience with offline reinforcement learning.  _arXiv preprint arXiv:2010.14500_ , 2020.
  * Sutton & Barto (2018) Richard S Sutton and Andrew G Barto.  _Reinforcement learning: An introduction_.  Second edition, 2018.
  * Sutton et al. (2016) Richard S. Sutton, A. Rupam Mahmood, and Martha White.  An emphatic approach to the problem of off-policy temporal-difference learning.  _J. Mach. Learn. Res._ , 17(1):2603–2631, January 2016.  ISSN 1532-4435.
  * Tian et al. (2020) Yuandong Tian, Lantao Yu, Xinlei Chen, and Surya Ganguli.  Understanding self-supervised learning with dual deep networks.  _arXiv preprint arXiv:2010.00578_ , 2020.
  * Tian et al. (2021) Yuandong Tian, Xinlei Chen, and Surya Ganguli.  Understanding self-supervised learning dynamics without contrastive pairs.  _arXiv preprint arXiv:2102.06810_ , 2021.
  * van Hasselt et al. (2018) Hado van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil.  Deep reinforcement learning and the deadly triad.  _ArXiv_ , abs/1812.02648, 2018.
  * Van Hasselt et al. (2018) Hado Van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil.  Deep reinforcement learning and the deadly triad.  _arXiv preprint arXiv:1812.02648_ , 2018.
  * Wang et al. (2021a) Ruosong Wang, Dean Foster, and Sham M. Kakade.  What are the statistical limits of offline {rl} with linear function approximation?  In _International Conference on Learning Representations_ , 2021a.  URL https://openreview.net/forum?id=30EvkP2aQLD.
  * Wang et al. (2021b) Ruosong Wang, Yifan Wu, Ruslan Salakhutdinov, and Sham M Kakade.  Instabilities of offline rl with pre-trained neural representation.  _arXiv preprint arXiv:2103.04947_ , 2021b.
  * Wei et al. (2019) Colin Wei, Jason Lee, Qiang Liu, and Tengyu Ma.  Regularization matters: Generalization and optimization of neural nets vs their induced kernel.  2019\.
  * Woodworth et al. (2020) Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro.  Kernel and rich regimes in overparametrized models.  _arXiv preprint arXiv:2002.09277_ , 2020.
  * Wu et al. (2019) Yifan Wu, George Tucker, and Ofir Nachum.  Behavior regularized offline reinforcement learning.  _arXiv preprint arXiv:1911.11361_ , 2019.
  * Xie & Jiang (2020) Tengyang Xie and Nan Jiang.  Q* approximation schemes for batch reinforcement learning: A eoretical comparison.  2020\.
  * Zanette (2020) Andrea Zanette.  Exponential lower bounds for batch reinforcement learning: Batch rl can be exponentially harder than online rl.  _arXiv preprint arXiv:2012.08005_ , 2020.
  * Zhang et al. (2019) Chiyuan Zhang, Samy Bengio, and Yoram Singer.  Are all layers created equal?  _arXiv preprint arXiv:1902.01996_ , 2019.

## Appendices

###  Appendix A Additional Visualizations and Experiments for DR3

In this section, we provide visualizations and diagnostic experiments
evaluating various aspects of feature co-adaptation and the DR3 regularizer.
We first provide more empirical evidence showing the presence of feature co-
adaptation in modern deep offline RL algorithms. We will also visualize DR3
inspired from the implicit regularizer term in TD-learning alleviates rank
collapse discussed in Kumar et al. (2021). We will compare the efficacies of
the explicit regularizer induced for different choices of the noise covariance
matrix $M$ (Equation 3.1), understand the effect of dropping the stop gradient
term in ou practical regularizer and finally, perform diagnostic experiments
visualizing if the Q-networks learned with DR3 resemble more like neural
networks trained via supervised learning, measured in terms of sensitivity and
robustness to layer reinitialization (Zhang et al., 2019).

####  A.1 More Empirical Evidence of Feature Co-Adaptation

In this section, we provide more empirical evidence demonstrating the
existence of the feature co-adaptation issue in modern offline RL algorithms
such as DQN and CQL. As shown below in Figure A.1, while the average dataset
Q-value for both CQL and DQN exhibit a flatline trend, the dot product
similarity for consecutive state-action tuples generally continues to increase
throughout training and does not flatline. While DQN eventually diverges in
Seaquest, the dot products increase with more gradient steps even before
divergence starts to appear.

Figure A.1:  Demonstrating feature co-adaptation on five Atari games with
standard offline DQN and CQL, averaged over 3 seeds. Observe that the feature
dot products continue to rise with more training for both CQL and DQN,
indicating the presence of co-adaptation. On the other hand, the average
Q-values exhibit a converged trend, except on Seaquest. Further, note that the
dot products continue to increase for CQL even though CQL explicitly corrects
for out-of-distribution action inputs.

####  A.2 Layer-wise structure of a Q-network trained with DR3

To understand if DR3 indeed makes Q-networks behave as if they were trained
via supervised learning, utilizing the empirical analysis tools from Zhang et
al. (2019), we test the robustness/sensitivity of each layer in the learned
network to re-initialization, while keeping the other layers fixed. This tests
if a particular layer is _critical_ to the predictions of the learned neural
network and enables us to reason about generalization properties (Zhang et
al., 2019; Chatterji et al., 2019). We ran CQL and REM and saved all the
intermediate checkpoints. Then, as shown in Figure A.2, we first loaded a
checkpoint ($x$-axis), and computed policy performance (shaded color;
colorbar) by re-initializing a given layer ($y$-axis) of the network to its
initialization value before training for the same run.

Figure A.2:  CQL vs CQL + DR3 and REM vs REM + DR3. Average robustness of the
learned Q-function to re-initialization of all layers to different checkpoints
over the course of training created based on the protocol from Zhang et al.
(2019). The colors in the heatmap indicate performance of the reinitialized
checkpoint, normalized w.r.t. the checkpoint without any change to layers.
Note that while CQL and REM are more sensitive (i.e., less robust) to
reinitialization of all the layers especially the last layer, CQL + DR3 and
REM + DR3 behave closer to supervised learning, in the sense that they are
more robust to reinitialization of layers of the network, especially the last
layer.

Note in Figure A.2, that while almost all layers are absolutely critical for
the base CQL algorithm, utilizing DR3 substantially reduces sensitivity to the
latter layers in the Q-network over the course of training. This is similar to
what Zhang et al. (2019) observed for supervised learning, where the initial
layers of a network were the most critical, and the latter layers primarily
performed near-random transformations without affecting the performance of the
network. This indicates that utilizing DR3 alters the internal layers of a
Q-network trained with TD to behave closer to supervised learning.

####  A.3 Results on MuJoCo Domains

In this section, we provide the results of applying DR3 on the MuJoCo tasks
shown in Figure A.5. (Appendix A) of Kumar et al. (2021). To briefly describe
the setup, in these tasks we train on the three gym tasks (Hopper-v2, Ant-v2,
Walker2d-v2) using 20% of the offline data, uniformly subsampled from the run
of an online SAC agent, mimicking the setup from Kumar et al. (2021). Rather
than retraining an SAC agent to collect data, we subsampled the Gym-MuJoCo
*-full-replay-v2 replay buffers from the latest D4RL (Fu et al., 2020). In
these cases we plot the $\mathrm{srank}_{\delta}$ values, the feature dot
products and the corresponding performance values with and without the DR3
regularizer for 4M steps (Kumar et al. (2021) showed their plots for just
under 4M steps) in Figure A.3.

Figure A.3:  Comparison of CQL and CQL + DR3 on the offline MuJoCo Gym
domains, mimicking the setup of Kumar et al. (2021). The data is generated by
randomly sampling 20% of the transitions of the D4RL (Fu et al., 2020) full-
replay-v2 datasets, which are collected via the run of an online SAC agent.
The performance is shown in terms of the D4RL normalized score, where 0.0
denotes the performance of a random policy and 100.0 denotes the performance
of an expert online SAC policy. Observe that adding DR3 stabilizes the
performance on Hopper, and prevents performance collapse on Walker2d and Ant.
In addition note that the srank values attained by CQL + DR3 is higher than
base CQL and more importantly, the feature dot products are much smaller for
CQL + DR3 compared to CQL.

Observe in Figure A.3, that while the standard CQL algorithm performs poorly
and suffers from performance degradation within about 1M-1.5M steps for
Walker2d and Ant, CQL + DR3 is able to prevent the performance degradation and
trains stably. Base CQL demonstrates oscillatory performance on Hopper, but
CQL + DR3 stabilizes the performace of CQL. This indicates that DR3 is
effective on MuJoCo domains, and prevents the instabilities with CQL.

For details, the weight on the CQL regularizer in this case is equal to 5.0
across all the tasks, and weight on the DR3 regularizer is 0.01. We also
attempted to tune the CQL coefficient for the baseline CQL algorithm within
$\\{1.0,2.0,5.0,10.0,20.0\\}$ to see if it address the performance degradation
issues, but did not find any difference in the collapsing behavior of base
CQL. Our CQL baseline is therefore well-tuned, and DR3 improves the
performance over this baseline.

####  A.4 Rank Collapse is Alleviated With DR3

Figure A.4:  Comparing the feature ranks for CQL and CQL + DR3. Observe that
utilizing DR3 successfully alleviates the rank collapse issue noted in prior
work without explicitly correcting for it.

Prior work (Kumar et al., 2021) has shown that implicit regularization in TD-
learning can lead to a feature rank collapse phenomenon in the Q-function,
which hinders the Q-function from using its full representational capacity.
Such a phenomenon is absent in supervised learning, where the feature rank
does not collapse. Since DR3 is inspired by mitigating the effects of the term
in the implicit regularizer (Equation 3.1) that only appears in the case of
TD-learning, we wish to understand if utilizing DR3 also alleviates rank
collapse. To do so, we compute the effective rank
$\mathrm{srank}_{\delta}(\phi)$ metric of the features learned by Q-functions
trained via CQL and CQL with DR3 explicit regularizer. As shown in Figure A.4,
for the case of five Atari games, utilizing DR3 alleviates the rank collapse
issue completely (i.e., the ranks do not collapse to very small values when
CQL + DR3 is trained for long). We do not claim that the ranks with DR3 are
necessarily higher, and infact as we show below, a higher srank of features
may not always imply a better solution. The fact that DR3 can prevent rank
collapse is potentially surprising, because no term in the practical DR3
regularizer explicitly aims to increase rank: feature dot products can be made
smaller while retaining low ranks by simply rescaling the feature vectors.
But, as we observe, utilizing DR3 enables learning features that do not
exhibit collapsed ranks, thus we hypothesize that correcting for appropriate
terms in $\mathrm{R}_{\mathrm{TD}}(\theta)$ can address some of the previously
observed pathologies in TD-learning.

Figure A.5:  Performance and $\mathrm{srank}$ values for DQN and DQN + DR3.
Observe that the srank values increase for DQN + DR3, while they collapse for
DQN on Asterix, Seaquest and SpaceInvaders with more training. Thus, DQN + DR3
does not suffer from a sudden rank collapse. However, a higher srank does not
imply a better return, and so while initially DQN does have a high rank, DQN +
DR3 performs superiorly.

We now investigate the feature ranks of a Q-network trained when DR3 is
applied in conjunction with a standard DQN and REM (Agarwal et al., 2020) on
the Atari domains. We plot the values of $\mathrm{srank}_{\delta}(\phi)$, the
feature dot products and the performance of the algorithm for DQN in Figure
A.5 and for REM in Figure A.6. In the case of DQN, we find that unlike the
base DQN algorithm for which feature rank does begin to collapse with more
training, the srank for DQN + DR3 is increasing. We also note that DQN + DR3
attains a better performance compared to DQN, throughout training.

Figure A.6:  Comparing the performance and $\mathrm{srank}$ values for REM and
REM + DR3. Observe that while REM + DR3 outperforms REM, the $\mathrm{srank}$
values attained by REM are much larger than REM + DR3, and none of these ranks
have collapsed. Thus, while REM + DR3 maintains non-collapsed features, for
the case of REM, it reduces the value of $\mathrm{srank}$ and attains better
performance. This does not contradict the observations from Kumar et al.
(2021) as we discuss in the text.

However, we note that the opposite trend is true for the case of REM: while
REM + DR3 attains a better performance than REM, adding DR3 leads to a
reduction in the $\mathrm{srank}$ value compared to base REM. At a first
glance, this might seem contradicting Kumar et al. (2021), but this is not the
case: to our understanding, Kumar et al. (2021) establish a correlation
between extremely low rank values (i.e., rank collapse) and poor performance,
but this does not mean that all high rank features will lead to good
performance. We suspect that since REM trains a multi-headed Q-function with
shared features and randomized target values, it is able to preserve high-rank
features, but this need not mean that these features are useful. In fact, as
shown in Figure A.7, we find that the base REM algorithm does exhibit feature
co-adaptation. This case is an example where the srank metric from Kumar et
al. (2021) may not indicate poor performance.

Figure A.7:  Feature dot products for REM and REM + DR3 on log scale. REM does
suffer from feature co-adaptation despite high-rank features.

####  A.5 Induced Implicit Regularizer: Theory And Practice

Table A.1: Normalized interquartile mean performance with 95% stratified bootstrap CIs (Agarwal et al., 2021) across 17 Atari games of REM, REM + $\Delta^{\prime}(\Phi)$ (Stop gradient in DR3), REM + DR3 after 6.5M gradient steps for the 1% setting and 12.5M gradient steps for the 5%, 10% settings. Observe that REM + $\Delta^{\prime}(\phi)$ also improves over the base REM method significantly, by about 130%, even though $\Delta^{\prime}(\phi)$ is generally comparable and somewhat worse than the DR3 regularizer used in the main paper. Data | REM | REM + $\Delta^{\prime}(\Phi)$ | REM+DR3
---|---|---|---
1% | 4.0 (3.3, 4.8) | 15.0 (13.4, 16.6) | 16.5 (14.5, 18.6)
5% | 25.9 (23.4, 28.8) | 55.5 (50.8, 59.8) | 60.2 (55.8, 65.1)
10% | 53.3 (51.4, 55.3) | 67.7 (64.7, 71.3) | 73.8 (69.3, 78)

In this section, we compare the performance of our practical DR3 regularizer
to the regularizers (Equation 3.1) obtained for different choices of $M$, such
as $M$ induced by noise, studied in previous work, and also evaluate the
effect of dropping the stop gradient function from the practical version of
our regularizer.

Empirically comparing the explicit regularizers for different noise covariance
matrices, $M$. The theoretically derived regularizer (Equation 3.1) suggests
that for a given choice of $M$, the following equivalent of feature dot
products should increase over the course of training:

$\Delta_{M}(\theta):=\sum_{\mathbf{s},\mathbf{a}\in\mathcal{D}}\mathrm{trace}\left[\Sigma^{*}_{M}\nabla
Q(\mathbf{s},\mathbf{a})\nabla
Q(\mathbf{s}^{\prime},\mathbf{a}^{\prime})^{\top}\right].~{}~{}~{}~{}~{}~{}\text{(Generalized
dot products)}$ (A.1)

We evaluate the efficacy of the explicit regularizer that penalizes the
generalized dot products, $\Delta_{M}(\theta)$, in improving the performance
of the policy, with the goal of identifying if our practical method performs
similar to this regularizer on generalized dot products.. While
$\Sigma^{*}_{M}$ must be explicitly computed by running fixed point iteration
for every parameter iterate $\theta$ found during TD-learning – which makes
this method significantly computationally expensive222In our implementation,
we run 20 steps of the fixed-point computation of $\Sigma$ as shown in Theorem
3.1 for each gradient step on the Q-function, and this increases the runtime
to about 8 days for 50 iterations on a P100 GPU., we evaluated it on five
Atari games for 50 $\times$ 62.5k gradient steps as a proof of concept. As
shown in Figure A.8, the DR3 penalty with the choice of $M$ which corresponds
to label noise, and the dot product DR3 penalty, which is our main practical
approach in this paper generally perform similarly on these domains, attaining
almost identical learning curves on 4/5 games, and clearly improving over the
base algorithm. This hints at the possibility of utilizing other noise
covariance matrices to derive an explicit regularizer. Deriving more
computationally efficient versions of the regularizer for a general $M$ and
identifying the best choice of $M$ are subject to future work.

Figure A.8:  Comparing the performance of explicit penalties for two different
choices of the covariance matrix $M$. Observe that in all the five games the
DR3 regularizer derived for the choice of $M$ from Blanc et al. (2020) also
leads to a substantial increase in performance over the base algorithm, and in
four of five games, DR3 (label-noise) works just as well as DR3.

Effect of stop gradient. Finally, we investigate the effect of utilizing a
stop gradient in the DR3 regularizer. We run a variant of DR3:
$\Delta^{\prime}(\phi)=\sum_{\mathbf{s},\mathbf{a},\mathbf{s}^{\prime}}\phi(\mathbf{s},\mathbf{a})^{\top}[[\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})]]$,
with the stop gradient on the second term
$(\mathbf{s}^{\prime},\mathbf{a}^{\prime})$ and present a comparison to the
one without the stop gradient in Table A.1 for REM as the base offline method,
averaged over 17 games. Note that this version of DR3, with the stop gradient,
also improves upon the baseline offline RL method (i.e., REM) by 130%. While
this performs largely similar, but somewhat worse than the complete version
without the stop gradient, these results do indicate that utilizing
$\Delta^{\prime}(\phi)$ can also lead to significant gains in performance.

####  A.6 Understanding Feature Co-Adaptation Some More

In this section, we present some more empirical evidence to understand feature
co-adaptation. The three factors we wish to study are: (1) the effect of
target update frequency on feature co-adaptation; (2) understand the trend in
normalized similarities and compare these to the trend in dot products; and
(3) understand the effect of out-on-sample actions in TD-learning and compare
it to offline SARSA on a simpler gridworld domain. We answer these questions
one by one via experiments aiming to verify each hypothesis.

Figure A.9:  Comparing the feature dot products for various target update
delays, where a smaller $N$ implies a faster update and a larger $N$
corresponds to a slower target update. Observe that while slower updates to
the target network may reduce co-adaptation, very slow target updates may
still lead to excesssive co-adaptation.

#####  A.6.1 Effect of Target Update Frequency on Feature Co-Adaptation

We studied the effect of target update frequency on feature co-adaptation, on
some gridworld domains from Fu et al. (2019). We utilized the grid16smoothobs
environment, where the goal of the agent is to navigate from the center of a
$16\times 16$ gridworld maze to one of its corners while avoiding obstacles
and “lava” cells. The observations provided to the RL algorithm are given by a
high-dimensional random transformation of the $(x,y)$ coordinates, smoothed
over neighboring cells in the gridworld. We sampled an offline dataset of 256
transitions and trained a Q-network with two hidden layers of size
$(1024,1024)$ via fitted Q-iteration (FQI) (Riedmiller, 2005).

We evaluated the feature dot products for Q-functions trained with a varying
target update frequencies, given generically as: updating the target network
using a hard target update once per $N$ gradient steps, where $N$ takes on
values $N=5,10,50,100,200,500$, and present the results in Figure A.9 (left),
averaged over 3 random seeds. Thee feature dot products initially decrease
from $N=5$ to $N=10$, because the target network is updated slower, but then
starts to rapidly increase when when the target network is slowed down further
to $N=50$ and $N=200$ in one case and $N=500$ in the other case.

We also evaluated the feature dot products when using CQL as the base offline
RL algorithm. As shown in Figure A.9 (right), while CQL does reduce the
absolute range of the feature dot products, slow target updates with $N=500$
still lead to the highest feature dot products as training progresses.
Takeaway: While it is intuitive to think that a slower target network might
alleviate co-adaptation, we see that this is not the case empirically with
both FQI and CQL, suggesting a deeper question that is an interesting avenue
for future study.

#####  A.6.2 Gridworld Experiments Comparing TD-learning vs Offline SARSA

To supplement the analysis in Section 3.1, we ran some experiments in the
gridworld domains from Fu et al. (2019). In this case, we used the
grid16smoothsparse and grid16randomsparse domains, which present challenging
navigation tasks in a maze under a 0-1 sparse reward signal, provided at the
end of the trajectory. Additionally, the observations available to the offline
RL agent do not consist of the raw $(x,y)$ locations of the agent in the maze,
but rather high-dimensional randomly chosen transformations of $(x,y)$ in the
case of grid16randomsparse, which are additionally smoothed locally around a
particular state to obtain grid16smoothsparse.

Since our goal is to compare feature co-adaptation in TD-learning and offline
SARSA, we consider a case where we evaluate a “mixed” behavior policy that
chooses the optimal action with a probability of 0.7 at a given state, and
chooses a random, suboptimal action with 0.3. We then generate a dataset of
size $256$ transitions and train offline SARSA and TD-learning on this data.
While SARSA backups the next action observed in the offline dataset, TD-
learning computes a full expectation of the Q-function
$\mathbb{E}_{\mathbf{a}^{\prime}\sim\pi_{\beta}(\cdot|\mathbf{s}^{\prime})}\left[Q(\mathbf{s}^{\prime},\mathbf{a}^{\prime})\right]$
under the behavior policy for computing Bellman backup targets. The behavior
policy is fully known to the TD-learning agent. Our Q-network consists of two
hidden layers of size $(1024,1024)$ as before.

Figure A.10:  Comparing the feature dot products for TD-learning and offline
SARSA, used to compute the value of the behavior policy using a dataset of
size 256 on two gridworld domains. Observe that the feature dot products are
higher in the case of TD-learning compared to offline SARSA.

We present the trends in the feature dot products for TD-learning and offline
SARSA in Figure A.10, averaged over three seeds. Observe that the trends in
the dot product values for TD-learning and offline SARSA closely follow each
other for the initial few gradient steps, soon, the dot products in TD-
learning start growing faster. In contrast, the dot products for SARSA either
saturate or start decreasing. The only difference between TD-learning and
SARSA is the set of actions used to compute Bellman targets – while the
actions used for computing Bellman backup targets in SARSA are in-sample
actions and are observed in the dataset, the actions used by TD-learning may
be out-of-sample, but are still within the distribution of the data-generating
behavior policy. This supports our empirical evidence in the main paper
showing that out-of-sample actions can lead to feature co-adaptation.

#####  A.6.3 Feature Co-Adaptation and Normalized Feature Similarities

Note that we characterized feature co-adaptation via the dot products of
features. In this section, we explore the trends in other notions of
similarity, such as cosine similarity between $\phi(\mathbf{s},\mathbf{a})$
and $\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})$ which measures the dot
product of feature vectors at consecutive state-action tuples after
normalization. Formally,

$\displaystyle\text{cos}(\phi(\mathbf{s},\mathbf{a}),\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})):=\frac{\phi(\mathbf{s},\mathbf{a})^{\top}\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})}{||\phi(\mathbf{s},\mathbf{a})||_{2}\cdot||\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})||_{2}}.$

We plot the trend in the cosine similarity with and without DR3 for five
Atari games in Figure A.11 with CQL, DQN and REM, and for the three MuJoCo
tasks studied in Appendix A.3 in Figure A.12. We find that the cosine
similarity is generally very high on the Atari domains, close to 1, and not
indicative of performance degradation. On the Ant and Walker2d MuJoCo domains,
we find that the cosine similarity first rises up close to 1 and roughly
saturates there. On the Hopper domain, the cosine similarity even decreases
over training. However we observe that the feature dot products are increasing
for all the domains. Applying DR3 in both cases improves performance (as shown
in earlier Appendix), and generally gives rise to reduced cosine similarity
values, though it can also increase the cosine similarity values occasionally.
Furthermore, even when the cosine similarities were decreasing for the base
algorithm (e.g., in the case of Hopper), addition of DR3 reduced the feature
dot products and helped improve performance. This indicates that both the norm
and directional alignment are contributors to the co-adaptation issue, which
is what DR3 aims to fix and independently directional alignment does not
indicate poor performance.

Figure A.11:  Cosine similarities of DQN, DQN + DR3, REM, REM + DR3 and CQL,
CQL + DR3. Note that DQN, REM and CQL attain close to 1 cosine similarities,
and addition of DR3 does reduce the cosine similarities of consecutive state-
action features. Figure A.12:  Cosine similarities of CQL and CQL + DR3 on
MuJoCo domains. Note that the cosine similarities of CQL grow to 1 and roughly
stabilize for Ant and Walker2d, but start decreasing for Hopper. This happens
despite the oscillatory trends in performance of CQL on Hopper A.3. This means
that a low cosine similarity need not imply poor performance, and DR3 can
improve performance even when cosine similarity of base CQL is decreasing. We
also notice that DR3 does actually reduce cosine similarity.

####  A.7 Stability of DR3 From a Good Solution

In this appendix, we study the trend of CQL + DR3 when starting learning from
a good initialization, which was studied in Figure 2. As shown in Figure A.13,
while the performance for baseline CQL degrades significantly (from 5000 at
initialization on Asterix, performance degrades to $\sim$2000 by 100
iterations for base CQL), whereas the performance of DR3 only moves from 5000
to $\sim$4300\. A similar trend holds for Breakout. This means that the
addition of DR3 does stabilize the learning relative to the baseline
algorithm. Please note that we are not claiming that DR3 is unequivocally
stable, but that improves stability relative to the base method.

Figure A.13:  Running CQL + DR3 and CQL in the setup of Figure 2 to evaluate
the stability of CQL + DR3 when starting training from a good solution.
Observe that the performance of base CQL decays quickly from the good
solution, but CQL + DR3 is _relatively_ more stable. Figure A.14:  Running DQN
+ DR3 and DQN in the setup of Figure 2 to evaluate the stability of DQN + DR3
when starting training from a good solution. Observe that the performance of
base DQN decays quickly from the good solution, but DQN + DR3 is _relatively_
more stable.

####  A.8 Statistical Significance of DR3 and Franka Kitchen Results

Table A.2: Performance of CQL, CQL + DR3 after 2M gradient steps on the Franka Kitchen domains averaged over 3 seeds. This is training for 6x longer compared to CQL defaults. Observe that CQL + DR3 outperforms CQL at 2M steps, indicating is efficacy in preventing long term performance degradation.  D4RL Task | CQL | CQL + DR3
---|---|---
kitchen-mixed | 27.67 $\pm$ 12.66 | 37.00 $\pm$ 11.53
kitchen-partial | 20.67 $\pm$ 15.57 | 40.67 $\pm$ 4.04
kitchen-complete | 28.00 $\pm$ 14.73 | 38.67 $\pm$ 6.66

We present the results comparing CQL and CQL+DR3 on the Franka Kitchen tasks
from D4RL in Table A.2. Observe that CQL+DR3 outperforms CQL, and to test the
statistical significance of these results, we analyze the probability of
improvement of CQL+DR3 over CQL next.

Figure A.15:  Statistical significance of the results of CQL + DR3 vs CQL
(Table 2) as measured by average probability of improvement (Agarwal et al.,
2021), with stratified bootstrap confidence intervals for this statistic.
Since the lower CI for this statistic is $>0.5$, CQL + DR3 significantly
improves over base CQL, and since the mean and upper CI are $\geq 0.75$, this
improvement is also meaningful.

In order to assess the statistical significance of our D4RL Antmaze and
Kitchen results, we follow the recommendedations by Agarwal et al. (2021) for
comparing deep RL algorithms considering their statistical uncertainties.
Specifically, we computed the average probability of improvement (Agarwal et
al., 2021) of CQL + DR3 over CQL on the antmaze and kitchen domains, and we
find that DR3 does significantly and meaningfully improve over CQL on both the
Kitchen and AntMaze domains. Before presenting the results, let us first
describe the metric we compute.

Probability of improvement and statistical significance. For two given
algorithms $\mathsf{Alg}_{1}$ and $\mathsf{Alg}_{2}$, and runs
$X_{k,1},X_{k,2},\cdots,X_{k,m}$ from $\mathsf{Alg}_{1}$ and runs
$Y_{k,1},Y_{k,2},\cdots,Y_{k,n}$ from $\mathsf{Alg}_{2}$ on task $k$, the
probability of improvement of $\mathsf{Alg}_{1}$ over $\mathsf{Alg}_{2}$ is
given by
$P(\mathsf{Alg}_{1}>\mathsf{Alg}_{2})=\frac{1}{K}\sum_{k=1}^{K}P(\mathsf{Alg}^{k}_{1}>\mathsf{Alg}^{k}_{2})$.
The probability of improvement on a given task $k$,
$P(\mathsf{Alg}^{k}_{1}>\mathsf{Alg}^{k}_{2})$ is computed using the Mann-
Whitney U-statistic and is given by:

$\displaystyle
P(\mathsf{Alg}^{k}_{1}>\mathsf{Alg}^{k}_{2})=\frac{1}{MN}\sum_{i=1}^{M}\sum_{j=1}^{N}S(X_{k,i},Y_{k,j})\quad\text{where}\quad
S(x,y)={\begin{cases}1,&{\text{if }}y<x,\\\ {\tfrac{1}{2}},&{\text{if
}}y=x,\\\ 0,&{\text{if }}y>x.\end{cases}}$

$\mathsf{Alg}_{1}$ leads to statistically _significant_ improve over
$\mathsf{Alg}_{2}$ if the lower CI for
$\mathsf{P}(\mathsf{Alg}_{1}>\mathsf{Alg}_{2})$ is larger than 0.5. Per the
Neyman-Pearson statistical testing criterion in Bouthillier et al. (2021),
$\mathsf{Alg}_{1}$ leads to statistically _meaningful_ improvement over
$\mathsf{Alg}_{2}$ if the upper confidence interval (CI) of
$\mathsf{P}(\mathsf{Alg}_{1}>\mathsf{Alg}_{2})$ is larger than 0.75.

Figure A.15 presents the value of $\mathsf{P}(\text{CQL + DR3}>\text{CQL})$ on
the AntMaze and Kitchen domains at 2M gradient steps along with the 95% CI for
this statistic. DR3 improves over CQL on the AntMaze domains with probability
0.83 with 95% CI (0.7, 0.96) and on the Kitchen domains with probability 0.8
with 95% CI (0.6, 1.0). These values pass the criterion of being both
statistically significant and meaningful per the above definitions, implying
that DR3 does significantly and meaningfully improve upon CQL on these
domains.

###  Appendix B Extended Related Work

In this section, we briefly review some extended related works, and in
particular, try to connect feature co-adaptation and implicit regularization
to various interesting results pertaining to RL lower-bounds with function
approximation and self-supervised learning.

Lower-bounds for offline RL. Zanette (2020) identifies hard instances for
offline TD learning of linear value functions when the provided features are
“aliased”. Note that this work does not consider feature learning or implicit
regularization, but their hardness result relies heavily on the fact the given
linear features are aliased in a special sense. Aliased features utilized in
the hard instance inhibit learning along certain dimensions of the feature
space with TD-style updates, necessitating an exponential sample size for
near-accurate value estimation, even under strong coverage assumptions. A
combination of Zanette (2020)’s argument, which provides a hard instance given
aliased features, and our analysis, which studies the emergence of co-
adapted/similar features in the offline deep RL setting, could imply that the
co-adaptation can lead to failure modes from the hard instance, even on
standard Offline RL problems, when provided with limited data.

Connections to self-supervised learning (SSL). Several modern self-supervised
learning methods (Grill et al., 2020; Chen & He, 2020) can be viewwed as
utilizing some form of bootstrapping where different augmentations of the same
input ($\mathbf{x}+\text{Aug}_{1},\mathbf{x}+\text{Aug}_{2}$) serve as
consecutive state-action tuples that appear on two sides of the backup. If we
may extrapolate our reasoning of feature co-adaptation to this setting, it
would suggest that performing noisy updates on a self-supervised bootstrapping
loss will give us feature representations that are highly similar for
consecutive state-action tuples, i.e., the representations for
$\phi(\mathbf{x}+\text{Aug}_{1})^{\top}\phi(\mathbf{x}+\text{Aug}_{2})$ will
be high. Intuitively, an easy way for obtaining high feature dot products is
for $\phi(\cdot)$ to capture only that information in $\cdot$, which is
agnostic to data augmentation, thus giving rise to features that are invariant
to transformations. This aligns with what has been shown in self-supervised
learning (Tian et al., 2020; 2021). Another interesting point to note is that
while such an explanation would indicate that highly co-adapted features are
beneficial in SSL, such features can be adverse in value-based RL as discussed
in Section 3.

Preventing divergence in deep TD-learning. Finally, we discuss Achiam et al.
(2019) which proposes to pre-condition the TD-update using the inverse the
neural tangent kernel (Jacot et al., 2018) matrix so that the TD-update is
always a contraction, for every $\theta_{k}$ found during TD-learning.
Intuitively, this can be overly restrictive in several cases: we do not need
to ensure that TD always contracts, but that is eventually stabilizes at good
solution over long periods of running noisy TD updates, Our implicit
regularizer (Equation‘3.1) derives this condition, and our theoretically-
inspired DR3 regularizer shows that empirically, it suffices to penalize the
dot product similarity in practice.

###  Appendix C Proof of Theorem 3.1

In this section, we will derive our implicit regularizer
$R_{\mathrm{TD}}(\theta)$ that emerges when performing TD updates with a
stochastic noise model with covariance matrix $M$. We first introduce our
notation that we will use throughout the proof, then present our assumptions
and finally derive the regularizer. Our proof utilizes the analysis techniques
from Blanc et al. (2020) and Damian et al. (2021), which analyze label-noise
SGD for supervised learning, however key modifications need to be made to
their arguments to account for non-symmetric matrices that emerge in TD
learning. As a result, the form of the resulting regularizer is very
different. To keep the proof concise, we will appeal to lemmas from these
prior works which will allow us to bound certain concentration terms.

####  C.1 Notation

The noisy TD-learning update for training the Q-function is given by:

$\displaystyle\\!\\!\\!\theta_{k+1}=\theta_{k}-\eta\underbrace{\left(\sum_{i}\nabla_{\theta}Q(\mathbf{s}_{i},\mathbf{a}_{i})\left(Q_{\theta}(\mathbf{s}_{i},\mathbf{a}_{i})\\!-\\!(r_{i}\\!+\\!\gamma{Q}_{\theta}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i}))\right)\right)}_{:=g(\theta)}\\!+\\!\eta\varepsilon_{k},~{}~{}~{}~{}\varepsilon_{k}\sim\mathcal{N}(0,M)$
(C.1)

where $g(\theta)$ denotes the parameter update. Note that $g(\theta)$ is not a
full gradient of a scalar objective, but it is a form of a “pseudo”-gradient
or “semi”-gradient. Let $\varepsilon_{k}$ denote an i.i.d.random noise that is
added to each update. This noise is sampled from a zero-mean Gaussian random
variable with covariance matrix $M$, i.e., $\mathcal{N}(0,M)$.

Let $\theta^{*}$ denote a point in the parameter space such that in the
vicinity of $\theta^{*}$, $g(\theta)\leq\mathscr{C}$, for a small enough
$\mathscr{C}$. Let $G(\theta)$ denote the derivative of $g(\theta)$ w.r.t.
$\theta$: $G(\theta)=\nabla_{\theta}g(\theta)$ and let $\nabla G(\theta)$
denote the third-order tensor $\nabla^{2}_{\theta}g(\theta)$. For notation
clarity, let $G=G(\theta^{*}),\nabla G=\nabla G(\theta^{*})$. Let $e_{i}$
denote the signed TD error for a given transition
$(\mathbf{s}_{i},\mathbf{a}_{i},\mathbf{s}^{\prime}_{i})\in\mathcal{D}$ at
$\theta^{*}$:

$\displaystyle
e_{i}=Q_{\theta^{*}}(\mathbf{s}_{i},\mathbf{a}_{i})-(r_{i}+\gamma
Q_{\theta^{*}}(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})).$ (C.2)

Since $\theta^{*}$ is a fixed point of the training TD error, $e_{i}=0$.
Following Blanc et al. (2020), we will assume that the learning rate in
gradient descent, $\eta$, is small and we will ignore terms that scale as
$\mathcal{O}(\eta^{1+\delta})$, for $\delta>0$. Our proof will rely on using a
reference Ornstein-Uhlenbeck (OU) process which the TD parameter iterates will
be compared to. Let $\zeta_{k}$ denote the $k$-th iterate of an OU process,
which is defined as:

$\zeta_{k+1}=(I-\eta
G)\zeta_{k}+\eta\varepsilon_{k},~{}~{}~{}\varepsilon_{k}\sim\mathcal{N}(0,M)$
(C.3)

We will drop $\theta$ from $\nabla_{\theta}$ to indicate that the gradient is
being computed at $\theta^{*}$, and drop $(\mathbf{s}_{i},\mathbf{a}_{i})$
from $Q(\mathbf{s}_{i},\mathbf{a}_{i})$ and instead represent it as $Q_{i}$
for brevity; we will represent
$Q(\mathbf{s}^{\prime}_{i},\mathbf{a}^{\prime}_{i})$ as $Q^{\prime}_{i}$. We
assume that $\nabla^{2}Q_{i}$ is $\mathscr{L}_{2}$-Lipschitz and
$\nabla^{3}Q_{i}$ is $\mathscr{L}_{3}$-Lipschitz throughout the parameter
space $\Theta$.

####  C.2 Proof Strategy

For a given point $\theta^{*}$ to be an attractive fixed point of TD-learning,
our proof strategy would be to derive the condition under which it mimics a
given OU noise process, which as we will show stays close to the parameter
$\theta^{*}$. This condition would then be interpreted as the gradient of a
“induced” implicit regularizer. If the point $\theta^{*}$ is not a stationary
point of this regularizer, we will show that the movement $\theta$ is large
when running the noisy TD updates, indicating that the regularizer, atleast in
part guides the dynamics of TD-learning. To show this, we would write out the
gradient update, isolate some terms that will give rise to the implicit
regularizer, and bound the remaining terms using contraction and concentration
arguments. The contraction arguments largely follow prior work (though with
key exceptions in handling contraction with asymmetric and complex eigenvalue
matrices), while the form of the implicit regularizer is different. Finally,
we will interpret the resulting update over large timescales to show that
learning is indeed guided by the implicit regularizer.

####  C.3 Assumptions and Conditions

Next, we present some key assumptions we will need for the proof. Our first
assumption is that the matrix $G\in\mathbb{R}^{d\times d}$ is of maximal rank
possible, which is equal to the number of datapoints $n$ and $n\ll d$, the
dimensionality of the parameter space. Crucially, this assumption do not imply
that $G$ is of full rank – it cannot be, because we are in the
overparameterized regime.

######  Assumption A1 ($G$ spans an $n$-dimensional basis.).

Assume that the matrix $G$ spans $n$-possible directions in the parameter
space and hence, attains the maximal possible rank it can.

The second condition we require is that the matrices $\sum_{i}\nabla
Q_{i}\nabla Q_{i}^{\top}$ and $M$ share the same $n$-dimensional basis as
matrix $G$:

######  Assumption A2.

$\sum_{i}\nabla Q_{i}\nabla Q_{i}^{\top}$, $M$, and $G$ span identical
$n$-dimensional subspaces.

This is a technical condition that is required. If this condition is not met,
as we will show the learning dynamics of noisy TD will not be a contraction in
certain direction in the parameter space and TD-learning will not stabilize at
such a solution $\theta^{*}$. In fact, we will utilize a stronger version of
this statement for TD-learning to converge, and we will discuss this shortly.

####  C.4 Lemmas Used In The Proof

Next, we present some lemmas that would be useful for proving the theoretical
result.

######  Lemma C.1 (Expressions for the first and and second-order derivatives
of $g(\theta)$.).

The following definitions and expansions apply to our proof:

$\displaystyle G(\theta^{*})$
$\displaystyle=\sum_{i}\nabla^{2}Q_{i}e_{i}+\sum_{i}\nabla Q_{i}(\nabla
Q_{i}-\gamma\nabla Q^{\prime}_{i})^{\top}$ $\displaystyle\nabla
G(\theta^{*})[\mathbf{v},\mathbf{v}]$
$\displaystyle=2\sum_{i}\nabla^{2}Q_{i}\mathbf{v}\mathbf{v}^{\top}(\nabla
Q_{i}-\gamma\nabla
Q^{\prime}_{i})+\sum_{i}\mathrm{tr}\left((\nabla^{2}Q_{i}-\gamma\nabla^{2}Q^{\prime}_{i})\mathbf{v}\mathbf{v}^{\top}\right)\nabla
Q_{i}+\nabla^{3}Q_{i}e_{i}$

Lemma C.1 presents a decomposition of the matrix $G$ and the directional
derivative of the third order tensor $\nabla G[\mathbf{v},\mathbf{v}]$ in
directions $\mathbf{v}$ and $\mathbf{v}$, which will appear in the Taylor
expansion layer. Note that at $\theta^{*}$ since $e_{i}=0$, the first term in
$G(\theta^{*})$ and the third term in $\nabla
G(\theta^{*})[\mathbf{v},\mathbf{v}]$ vanish. Lemma C.2 derives a fixed-point
recursion for the covariance matrix of the total noise accumulated in the OU-
process with covariance matrix $M$ and this will appear in our proof.

######  Lemma C.2 (Covariance of the random noise process $\zeta_{k}$).

Let $\zeta_{k}$ denote the OU process satisfying: $\zeta_{k+1}=(I-\eta
G)\zeta_{k}+\eta\varepsilon_{k}$, where $\varepsilon_{k}\sim\mathcal{N}(0,M)$,
where $M\succcurlyeq 0$. Then, $\zeta_{k+1}\sim\mathcal{N}(0,\Sigma)$, where
$\Sigma$ satisfies the discrete Lyapunov equation:

$\Sigma^{*}_{M}=(I-\eta G)\Sigma^{*}_{M}(I-\eta G)^{\top}+\eta^{2}M.$

###### Proof.

For the OU process, $\zeta_{k+1}=(I-\eta G)\zeta_{k}+{\eta}\varepsilon_{k}$,
since $\varepsilon_{k}$ is a Gaussian random variable, by induction so is
$\zeta_{k+1}$, and therefore the covariance matrix of $\zeta_{k+1}$ is given
by:

$\Sigma_{k+1}:=(I-\eta G)\Sigma_{k}(I-\eta G^{\top})+\eta^{2}M.$ (C.4)

Solving for the fixed point for $\Sigma_{k}$ gives the desired expression. ∎

In our proofs, we will require the following contraction lemmas to tightly
bound the magnitude of some zero-mean terms that will appear in the noisy TD
update under certain scenarios. Unlike the analysis in Damian et al. (2021)
and Blanc et al. (2020) for supervised learning with label noise, where the
contraction terms like $(I-\eta G)^{k}G$ are bounded by
$\approx\frac{1}{k\eta}$ intuitively because $I-\eta G$ is a contraction in
the subspace spanned by matrix $G$. However, this is not true for TD-learning
directly since terms like $(I-\eta G)^{k}S$ appear for a different matrix $S$.
Therefore, TD-learning will diverge from $\theta^{*}$ unless matrices $G$ and
$M$ have their corresponding eigenvectors assigned to the top eigenvalues be
approximately “aligned”. We formalize this definition next, and then provide a
proof of the concentration guarantee.

######  Definition 1 ($(\omega,C_{0})$-alignment).

Given a positive semidefinite matrix $A$, let $A=U_{A}\Lambda_{A}U_{A}^{\top}$
denote its eigendecomposition. Without loss of generality assume that the
eiegenvalues are arranged in decreasing order, i.e., $\forall
i>j,\Lambda_{A}(i)\leq\Lambda_{A}(j)$. Given another matrix $B$, let
$B=U_{B}\Lambda_{B}U_{B}^{H}$ denote its complex eigendecomposition, where
eigenvalues in $\Lambda_{B}$ are arranged in decreasing order of their complex
magnitudes, i.e., $\forall i>j,|\Lambda_{B}(i)|\leq|\Lambda_{B}(j)|$. Then the
matrix pair $(A,B)$ is said to be $(\omega,C_{0})$-aligned if
$|U_{B}^{H}(i)U_{A}(i)|\leq\omega$ and if $\forall~{}i,\Lambda_{A}(i)\leq
C_{0}|\Lambda_{B}(i)|$ for a constant $C_{0}$.

If two matrices are $(\omega,C_{0})$-aligned, this means that the
corresponding eigenvectors when arranged in decreasing order of eigenvalue
magnitude roughly align with each other. This condition would be crucial while
deriving the implicit regularizer as it will quantify the rate of contraction
of certain terms that define the neighborhood that the iterates of noisy TD-
learning will lie in with high probability. We will operate in the setting
when the matrix $G$ and $\sum_{i}\nabla Q_{i}\nabla Q_{i}^{\top}$ are
$(\omega,C_{0})$-aligned with each other, and matrix $M$ and $G$ are also
$(\omega,C_{0})$-aligned (note that we can consider
$\omega^{\prime},C^{\prime}_{0})$, which will not change our bounds and
therefore we go for less notational clutter). Next we utilize this notion of
alignment to show a particular contraction bound that extends the weak
contraction bound in Damian et al. (2021).

######  Lemma C.3.

Assume we are given a matrix $G$ such that $|\lambda_{i}(I-\eta
G)|\leq\rho_{0}<1$ for all $\lambda_{i}$ such that $\lambda_{i}\neq 0$. Let
$G=U\Lambda U^{H}$ be the complex eigenvalue decomposition of $G$ (since
almost every matrix is complex-diagonalizable). For a positive semi-definite
matrix $S$ that is $(\omega,C_{0})$-aligned with $G$, if
$S=U_{S}\Lambda_{S}U_{S}^{\top}$ is its eigenvalue decomposition, the
following contraction bound holds:

$\displaystyle||(I-\eta G)^{k}S||=\mathcal{O}\left(\frac{\omega C_{0}}{\eta
k}\right)$

###### Proof.

To prove this statement, we can expand $(I-\eta G)$ using its eigenvalue
decomposition only in the subspace that is jointly shared by $G$ and $M$, and
then utilize the definition of $\omega$-alignment to bound the terms.

$\displaystyle||(I-\eta G)^{k}S||$ $\displaystyle=||(I-\eta U\Lambda
U^{H})^{k}U_{S}\Lambda_{S}U_{S}^{\top}||$ (C.5)
$\displaystyle=\left|\left|(UU^{H}-\eta
U\Lambda_{U}U^{H})^{k}U_{S}\Lambda_{S}U_{S}^{\top}\right|\right|$ (C.6)
$\displaystyle=\left|\left|U\left(I-\eta\Lambda\right)^{k}U^{H}U_{S}\Lambda_{S}U_{S}^{\top}\right|\right|$
(C.7)
$\displaystyle\leq\omega\cdot||\left(I-\eta\Lambda\right)^{k}||\cdot\Lambda_{S}$
(C.8) $\displaystyle\leq\omega\cdot
C_{0}\cdot\left(\max_{i}~{}~{}|1-\eta\Lambda(i)|^{k}|\Lambda(i)|\right)$ (C.9)

Now we need to solve for the inner maximization term. When $\Lambda(i)$ is not
complex for any $i$, the term above is $\lesssim 1/\eta k$ using the result
from Damian et al. (2021), but when $\Lambda(i)$ is complex, this bound can
only hold under certain conditions. To note when this quantity is bounded, we
expand $|1-\eta x|^{k}$ for some complex number
$x=r(\cos\theta+\iota\sin\theta)$:

$\displaystyle|1-\eta x|^{k}$ $\displaystyle=\left|\left(1-\eta
r\cos\theta\right)+\iota\eta r\sin\theta\right|$ (C.10)
$\displaystyle=\left[\sqrt{\left(1-\eta
r\cos\theta\right)^{2}+\eta^{2}r^{2}\sin^{2}\theta}\right]^{k}=\left(1+\eta^{2}r^{2}-2\eta
r\cos\theta\right)^{k/2}$ (C.11) $\displaystyle\implies|1-\eta x|^{k}|x|$
$\displaystyle=\left(1+\eta^{2}r^{2}-2\eta r\cos\theta\right)^{k/2}r$ (C.12)
$\displaystyle\lesssim\frac{1}{\eta
k}~{}~{}~{}~{}\text{if}~{}\eta\leq\min_{i}\frac{\mathrm{Re}(\Lambda(i))}{|\Lambda(i)|}~{}~{}~{}~{}~{}\text{and}~{}~{}~{}~{}\infty~{}~{}\text{otherwise}.$
(C.13)

Plugging back the above expression in the bound above completes the proof. ∎

The proof of Lemma C.3 indicates that unless the learning rate $\eta$ and the
matrix $G$ are such that the $|\lambda_{i}(I-\eta G)|\leq\rho<1$ in directions
spanned by matrix $S$, such an expression may not converge. This is expected
since the matrix $I-\eta G$ will not contract in directions of non-zero
eigenvalues if the real part $r\cos\theta$ is negative or zero. Additionally,
we note that under Definition 1, we can extend several weak-contraction bounds
from Damian et al. (2021) (Lemmas 9-14 in Damian et al. (2021)) to our
setting.

Next, Lemma C.4 shows that the OU noise iterates are bounded with high
probability when Definition 1 holds:

######  Lemma C.4 ($\zeta_{k}$ is bounded with high probability).

With probability atleast $1-\delta$ and under Definition 1, $||\zeta_{k}||\leq
n\omega\sqrt{\eta C_{0}}\log\frac{1}{\delta}=\mathcal{O}(\sqrt{\eta})$.

###### Proof.

To prove this lemma, we first bound the trace of the covariance matrix
$\Sigma_{k+1}$ and then apply high probability bounds on the Martingale norm
concentration. The trace of the covariance matrix $\Sigma_{k+1}$ can be
bounded as follows (all the equations below are restricted to the dimensions
of non-zero eigenvalues of $G$):

$\displaystyle\mathrm{tr}\left[\Sigma_{k+1}\right]$ $\displaystyle=\sum_{j\leq
k}\mathrm{tr}\left[(I-\eta G)^{j}M(I-\eta G^{\top})^{j}\right]$ (C.14)
$\displaystyle=\sum_{j\leq k}\mathrm{tr}\left[(UU^{H}-\eta U\Lambda
U^{H})^{j}M(UU^{H}-\eta U\Lambda U^{H})^{j}\right]$ (C.15)
$\displaystyle=\sum_{j\leq
k}\mathrm{tr}\left[U(I-\eta\Lambda)^{j}U^{H}U_{M}\Lambda_{M}U_{M}^{\top}U(I-\eta\Lambda)^{j}U^{H}\right]$
(C.16) $\displaystyle=\sum_{j\leq
k}n\omega^{2}C_{0}\mathrm{tr}\left[|I-\eta\Lambda|^{j}\cdot|\Lambda|\cdot|I-\eta\Lambda|^{j}\right]$
(C.17) $\displaystyle\leq n\omega^{2}C_{0}\sum_{j\leq
k}n\cdot\max_{\lambda}(|1-\eta\lambda|^{2j}\cdot|\lambda|)\leq{\eta
n^{2}C_{0}\omega^{2}}$ (C.18)

Now, we can apply Corollary 1 from Damian et al. (2021) to obtain a bound on
$||\zeta_{k}||$ as with high probability, atleast $1-\delta$,
$||\zeta_{k}||\leq\sqrt{2\mathrm{tr}(\Sigma)\log\frac{1}{\delta}}=n\omega\sqrt{\eta
C_{0}}\log\frac{1}{\delta}$. ∎

####  C.5 Main Proof of Theorem 3.1

In this section, we present the main proof of Theorem 3.1. The proof involves
two components: (1) the part where we derive the regularizer, and (2) bounding
additional terms via concentration inequalities. Part (1) is specific to TD-
learning, while a lot of the machinery for part (2) is directly taken from
prior work (Damian et al., 2021) and Blanc et al. (2020). We focus on part (1)
here.

Our strategy is to analyze the learning dynamics of noisy TD updates that
originate at $\theta^{*}$. In a small neighborhood around $\theta^{*}$, we can
expand the noisy TD update (Equation 3) using Taylor’s expansion around
$\theta^{*}$ which gives:

$\displaystyle\theta_{k+1}=\theta_{k}-\eta
g(\theta_{k})+\eta\varepsilon_{k},~{}~{}\varepsilon_{k}\sim\mathcal{N}(0,M)$
(C.19) $\displaystyle\implies$
$\displaystyle\theta_{k+1}=\theta_{k}-\eta\left(g+G(\theta_{k}-\theta^{*})-\frac{\eta}{2}G[\theta_{k}-\theta^{*},\theta_{k}-\theta^{*}]\right)+\eta\varepsilon_{k}+\mathcal{O}(\eta||\theta_{k}-\theta^{*}||^{3}).$
(C.20)

Denoting $\nu_{k}:=\theta_{k}-\theta^{*}$, using the fact that
$||g(\theta^{*})||\leq\mathscr{C}$, we find that $\nu_{k}$ can be written as:

$\displaystyle\nu_{k+1}$ $\displaystyle=(I-\eta
G)\nu_{k}+\varepsilon_{k}+\frac{\eta}{2}G[\nu_{k},\nu_{k}]+\mathcal{O}(\eta||\nu_{k}||^{3}+\eta\mathscr{C})$
(C.21)

Since the OU process $\zeta_{k}$ stays in the vicinity of the point
$\theta^{*}$, and follows a similar recursion to the one above, our goal would
be to design a regularizer so that Equation C.21 closely follows the OU
process. Thus, we would want to bound the difference between the variable
$\nu_{k}$ and the variable $\zeta_{k}$, denoted as $r_{k}$ to be within a
small neighborhood:

$\displaystyle r_{k+1}=\nu_{k+1}-\zeta_{k+1}=(I-\eta
G)\underbracket{(\nu_{k}-\zeta_{k})}_{r_{k}}+\frac{1}{2}G[\nu_{k},\nu_{k}]+\mathcal{O}(\eta||\nu_{k}||^{3}+\eta\mathscr{C}).$

We can write down an expression for $r_{k}$ summing over all the terms:

$r_{k+1}=-\underbracket{\frac{\eta}{2}\sum_{j\leq k}(I-\eta G)^{k-j}\nabla
G[\nu_{k},\nu_{k}]}_{\text{term (a)}}+\underbracket{\sum_{j\leq k}(I-\eta
G)^{j}\left[\mathcal{O}(\eta||\nu_{k}||^{3}+\eta\mathscr{C})\right]}_{\text{term
(b)}}.$ (C.22)

Term (a) in the above equation is the one that can induce a displacement in
$r_{k}$ as $k$ increases and would be used to derive the regularizer, whereas
term (b) primarily consists of terms that concentrate to $0$. We first analyze
term (a) and then we will analyze the concentration terms later.

To analyze term (a), note that the term $\nabla G[\nu_{k},\nu_{k}]$, by Lemma
C.1, only depends on $\nu_{k}$ via the covariance matrix
$\nu_{k}\nu_{k}^{\top}$. So we will partition this term into two terms: (i) a
term that utilizes the asymptotic covariance matrix of the OU process and (ii)
errors due to a finite $k$ and stochasticity that will concentrate.

$\displaystyle 2\times\text{(a)}~{}$ $\displaystyle=\eta\sum_{j\leq k}(I-\eta
G)^{k-j}\nabla G[\nu_{k},\nu_{k}]$ (C.23) $\displaystyle=\sum_{j\leq k}(I-\eta
G)^{k-j}\nabla G[\zeta^{*},\zeta^{*}]+\sum_{j\leq k}(I-\eta G)^{k-j}\nabla
G([\nu_{k},\nu_{k}]-[\zeta^{*},\zeta^{*}]),$ (C.24)

The first term is a “bias” term and doesn’t concentrate to $0$, and will give
rise to the regularizer. We can break this term using Lemma C.1 as:

$\displaystyle\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\\!\nabla
G[\zeta^{*},\zeta^{*}]=$ $\displaystyle
2\sum_{i}\nabla^{2}Q_{i}\Sigma^{*}_{M}(\nabla Q_{i}-\gamma\nabla
Q^{\prime}_{i})+\sum_{i}\mathrm{tr}\left[(\nabla^{2}Q_{i}-\gamma\nabla^{2}Q^{\prime}_{i})\Sigma^{*}_{M}\right]\nabla
Q_{i}$ (C.25)

The regularizer $R_{\mathrm{TD}}(\theta)$ is the function such that:

$\displaystyle\nabla_{\theta}R_{\mathrm{TD}}(\theta)$
$\displaystyle=\sum_{i}\nabla^{2}Q_{i}\Sigma^{*}_{M}(\nabla Q_{i}-\gamma\nabla
Q^{\prime}_{i})$ (C.26) $\displaystyle\implies R_{\mathrm{TD}}(\theta)$
$\displaystyle=\sum_{i}\nabla Q_{i}\Sigma^{*}_{M}\nabla
Q_{i}^{\top}-\gamma\sum_{i}\mathrm{trace}\left(\Sigma^{*}_{M}\nabla
Q_{i}[[\nabla Q^{\prime}_{i}]]^{\top}\right),$ (C.27)

where $[[\cdot]]$ denotes the stop gradient operator. If the point
$\theta^{*}$ is a stationary point of the regularizer
$R_{\mathrm{TD}}(\theta)$, then Equations C.26 and C.27 imply that the first
term of Equation C.25 must be 0. Therefore in this case to show that
$\theta^{*}$ is attractive, we need to show that the other terms in Equations
C.25, C.24 and term (b) in Equation C.22 concentrate around $0$ and are
bounded in magnitude. The remaining part of the proof shown in Appendix C.7
provides these details, but we first summarize the main takeaways in the proof
to conclude the argument.

####  C.6 Summary of the Argument

We will show how to concentrate terms in Equation C.26 besides the regularizer
largely following the techniques from prior work, but we first summarize the
entire proof. The overall update to the vector $r_{k}$ which measures the
displacement between the parameter vector $\theta_{k}-\theta^{*}$ and the OU-
process $\zeta_{k}$ can be written as follows, and it is governed by the
derivative of the implicit regularizer (modulo error terms):

$r_{k+1}=-\frac{\eta}{2}\sum_{j\leq k}(I-\eta
G)^{k-j}\nabla_{\theta}R_{\mathrm{TD}}(\theta^{*})+\mathcal{O}\left(\sqrt{\eta
t}\cdot\mathrm{poly}(\mathscr{C},\mathscr{L}_{2},\mathscr{L}_{3},\omega,C_{0})\right).$
(C.28)

An important detail to note here is that since the regularizer consists of
$\Sigma^{*}_{M}$ and the size of $\Sigma^{*}_{M}$ (i.e, its eigenvalues), as
shown in Lemma C.4 depends on one factor of $\eta$. So, effectively the first
term in Equation C.28 does depend on two factors of $\eta$. Using Equation
C.28, we can write the deviation between $\theta^{*}$ and $\theta_{k}$ as:

$\displaystyle\nu_{k+1}$ $\displaystyle=\zeta_{k+1}-\frac{\eta}{2}\sum_{j\leq
k}(I-\eta
G)^{k-j}\nabla_{\theta}R_{\mathrm{TD}}(\theta^{*})+\mathcal{O}\left(\sqrt{\eta
t}\cdot\mathrm{poly}(\mathscr{C},\mathscr{L}_{2},\mathscr{L}_{3},\omega,C_{0})\right).$
(C.29)

The OU process $\zeta_{k}$ converges to $\theta^{*}$ in the subspace spanned
by $G$, since the condition $\rho(I-\eta G)<1$ is active in this subspace (if
the condition that $\rho(I-\eta G)<1$ in the subspace spanned by $G$ is not
true, then as Ghosh & Bellemare (2020) show, TD can diverge). Now, given $G$
satisfies this spectral radius condition, $\zeta_{k}$ would converge to
$\theta^{*}$ within a timescale of $\mathcal{O}\left(\frac{1}{\eta}\right)$
within this subspace, which as Blanc et al. (2020) put it is the strength of
the “mean-reversion” term. On the remaining directions (note that $d\gg n$),
the dynamics is guided by the regularizer, although with a smaller weight of
$\eta^{2}$.

####  C.7 Additional Proof Details: Concentrating Other Terms

We first concentrate the terms in Equation C.25. The cumulative effect of the
second term in Equation C.25 is given by:

$\displaystyle\eta\sum_{j\leq k}(I-\eta G)^{j-k}\nabla
Q_{i}\mathrm{tr}\left[(\nabla^{2}Q_{i}-\gamma\nabla^{2}Q^{\prime}_{i})\Sigma^{*}_{M}\right]$
(C.30) $\displaystyle\leq\eta\sum_{j\leq k}(I-\eta G)^{j-k}\nabla
Q_{i}\cdot\mathcal{O}\left(\mathscr{L}_{2}(1+\gamma)\sigma\right)\leq\mathcal{O}\left(\eta\sqrt{\frac{k}{\eta}}\omega_{0}C_{0}\mathscr{L}_{2}(1+\gamma)\sigma\right),$
(C.31)

which follows from the fact that $\nabla^{2}Q_{i}$ is
$\mathscr{L}_{2}$-Lipschitz, and using Lemma C.3 for contracting the remaining
terms.

Next, we turn to concentrating the second term in Equation C.24. This term
corresponds to the contribution of difference between the empirical covariance
matrix $\nu_{k}\nu_{k}^{\top}$ and the asymptotic covariance matrix
$\zeta^{*}\zeta^{*\top}$. We expand this term below using the form of $G$ from
Lemma C.1, and bound it one by one.

$\displaystyle\sum_{j\leq k}(I-\eta G)^{k-j}\nabla
G([\nu_{k},\nu_{k}]-[\zeta^{*},\zeta^{*}])$ (C.32) $\displaystyle=\sum_{j\leq
k}\sum_{i}(I-\eta
G)^{k-j}\nabla^{2}Q_{i}\left(\nu_{k}\nu_{k}-\zeta^{*}\zeta^{*\top}\right)(\nabla
Q_{i}-\gamma\nabla Q^{\prime}_{i})+\mathcal{O}\left(\sqrt{\eta
k}\omega_{0}C_{0}\mathscr{L}_{2}(1+\gamma)\sigma\right)$ (C.33)

Now, we note that the term
$\Delta_{k+1}:=\nu_{k+1}\nu_{k+1}^{\top}-\zeta^{*}\zeta^{*\top}$ can itself be
written as a recursion:

$\displaystyle\Delta_{k+1}$ $\displaystyle=(I-\eta G)(\Delta_{k})(I-\eta
G)^{\top}+\underbracket{(I-\eta
G)\zeta_{k}\varepsilon^{\top}+\varepsilon\zeta_{k}^{\top}(I-\eta
G)^{\top}}_{A_{k}}+\underbracket{\varepsilon\varepsilon^{\top}-\eta
M}_{B_{k}}$ (C.34)

Expanding the term $\Delta_{k+1}$ in terms of a summation over $k$, and
plugging it into the expression from Equation C.33 we get

$\displaystyle\sum_{i}\sum_{j\leq k}(I-\eta G)^{k-j}$
$\displaystyle\nabla^{2}Q_{i}(I-\eta G)^{j}\Delta_{0}(I-\eta G^{\top})^{j}$
(C.35) $\displaystyle+\sum_{i}$ $\displaystyle~{}\sum_{j\leq k}\sum_{p\leq
j}(I-\eta G)^{k-j}\nabla^{2}Q_{i}(I-\eta G)^{j-p-1}(A_{p}+B_{p})(I-\eta
G^{\top})^{j-p-1}$

Now by noting that if $G$ and $\nabla Q_{i}$ are $(\omega,C_{0})$-aligned,
then so are $G^{\top}$ and $\nabla Q_{i}$, we can finish the proof by
repeating the calculations used by Damian et al. (2021) (Appendix B, Equations
67-73) to bound the terms in Equation C.35 by $\mathcal{O}(\sqrt{\eta k})$,
but with an additional factor of $\omega^{2}C_{0}^{2}$.

Term (b) in Equation C.22. When $\mathscr{C}$ is small enough, we can bound
the term (b) using $\mathcal{O}(\sqrt{\eta k})$, similar to Damian et al.
(2021).

###  Appendix D Proof of Proposition 3.2

In this section, we will prove Proposition 3.2. First, we refer to Proposition
3.1 in Ghosh & Bellemare (2020), which shows that TD-learning is stable and
converges if and only if the matrix
$M_{\phi}=\Phi^{\top}(\Phi-\gamma\Phi^{\prime})$ has eigenvalues with all
positive real entries. Now note that if,

$\displaystyle\sum_{\mathbf{s},\mathbf{a}}\phi(\mathbf{s},\mathbf{a})^{\top}\phi(\mathbf{s},\mathbf{a})$
$\displaystyle\leq\gamma\sum_{\mathbf{s},\mathbf{a},\mathbf{s}^{\prime}}\phi(\mathbf{s}^{\prime},\mathbf{a}^{\prime})^{\top}\phi(\mathbf{s},\mathbf{a})$
(D.1) $\displaystyle\implies\mathrm{trace}\left(\Phi^{\top}\Phi\right)$
$\displaystyle\leq\gamma\mathrm{trace}\left(\Phi^{\top}\Phi^{\prime}\right)$
(D.2)
$\displaystyle\implies\mathrm{trace}\left[\Phi^{\top}\left(\Phi-\gamma\Phi^{\prime}\right)\right]\leq
0.$ (D.3)

Since the trace of a real matrix is the sum of real components of eigenvalues,
if for a given matrix $M$, $\mathrm{trace}(M)\leq 0$, then there exists
atleast one eigenvalue $\lambda_{i}$ such that $\mathrm{Re}(\lambda_{i})\leq
0$. If $\lambda_{i}<0$, then the learning dynamics of TD would diverge, while
if $\lambda_{i}=0$ for all $i$, then learning will not contract towards the TD
fixed point. This concludes the proof of this result.

###  Appendix E Experimental Details of Applying DR3

In this section, we discuss the practical experimental details and
hyperparameters in applying our method, DR3 to various offline RL methods. We
first discuss an overview of the offline RL methods we considered in this
paper, and then provide a discussion of hyperparameters for DR3.

####  E.1 Background on Various Offline RL Algorithms

In this paper, we consider four base offline RL algorithms that we apply DR3
on. These methods are detailed below:

REM. Random ensemble mixture (Agarwal et al., 2020) is an uncertainty-based
offline RL algorithm uses multiple parameterized Q-functions to estimate the
Q-values. During the Bellman backup, REM computes a random convex combination
of the target Q-values and then trains the Q-function to match this randomized
target estimate. The randomized target value estimate provides a robust
estimate of target values, and delays unlearning and performance degradation
that we typically see with standard DQN-style algorithms in the offline
setting. For instantiating REM, we follow the instantiation provided by the
authors and instantiate a multi-headed Q-function with 200 heads, each of
which serves as an estimate of the target value. These multiple heads branch
off the last-but-one layer features of the base Q-network. The objective for
REM is given by:

$\\!\\!\\!\\!\\!\\!\\!\\!\\!\min_{\theta}\mathbb{E}_{\mathbf{s},\mathbf{a},r,\mathbf{s}^{\prime}\sim\mathcal{D}}\left[\mathbb{E}_{\alpha_{1},\dots,\alpha_{K}\sim{\Delta}}\left[\ell_{\lambda}\left(\sum_{k}\alpha_{k}Q_{\theta}^{k}(\mathbf{s},\mathbf{a})-r-\gamma\max_{\mathbf{a}^{\prime}}\sum_{k}\alpha_{k}Q_{\theta^{\prime}}^{k}(\mathbf{s}^{\prime},\mathbf{a}^{\prime})\right)\right]\right]$
(E.1)

where $l_{\lambda}$ denotes the Huber loss while $P_{\Delta}$ denotes the
probability distribution over the standard K-1 simplex.

CQL. Conservative Q-learning (Kumar et al., 2020b) is an offline RL algorithm
that learns a conservative value function such that the estimated performance
of the policy under this learned value function lower-bounds its true value.
CQL modifies the Q-function training to incorporate a term that minimizes the
overestimated Q-values in expectation, while maximizing the Q-values observed
in the dataset, in addition to standard TD error. This CQL regularizer is
typically multiplied by a coefficient $\alpha$, and we pick $\alpha=0.1$ for
all our Atari experiments following Kumar et al. (2021) and $\alpha=5.0$ for
all our kitchen and antmaze D4RL experiments. Using
$\overline{y}_{k}(\mathbf{s},\mathbf{a})$ to denote the target values computed
via the Bellman backup (we use actor-critic backup for D4RL experiments and
the $\max_{\mathbf{a}^{\prime}}$ backup for standard Q-learning in our Atari
experiments following Kumar et al. (2020b)), the objective for training CQL is
given by:

$\\!\\!{\min_{Q}\alpha\left(\mathbb{E}_{\mathbf{s}\sim\mathcal{D}}\left[\log\sum_{\mathbf{a}}\exp(Q(\mathbf{s},\mathbf{a}))\right]-\mathbb{E}_{\mathbf{s},\mathbf{a}\sim\mathcal{D}}\left[Q(\mathbf{s},\mathbf{a})\right]\right)\\!+\\!\frac{1}{2}\\!\mathbb{E}_{\mathbf{s},\mathbf{a},\mathbf{s}^{\prime}\sim\mathcal{D}}\left[\left(Q(\mathbf{s},\mathbf{a})-\overline{y}_{k}(\mathbf{s},\mathbf{a})\right)^{2}\right]}.$

The deep Q-network utilized by us is a ReLU network with four hidden layers of
size $(256,256,256,256)$ for the D4RL experiments, while for Atari we utilize
the standard convolutional neural network from Agarwal et al. (2020); Kumar et
al. (2021) with 3 convolutional layers borrowed from the nature DQN network
and then a hidden feedforward layer of size $512$.

BRAC. Behavior-regularized actor-critic (Wu et al., 2019) is a policy-
constraint based actor-critic offline RL algorithm which regularizes the
policy to stay close to the behavior policy $\pi_{\beta}$ to prevent the
selection of “out-of-distribution” actions. In addition, BRAC subtracts this
divergence estimate from the target Q-values when performing the backup, to
specifically penalize target values that come from out-of-distribution action
inputs at the next state ($\mathbf{s}^{\prime},\mathbf{a}^{\prime})$.

Q-function:
$\displaystyle~{}~{}\min_{\theta}~{}~{}\mathbb{E}_{\mathbf{s},\mathbf{a}\sim\mathcal{D}}\left[\left(r(\mathbf{s},\mathbf{a})+\gamma\mathbb{E}_{\mathbf{a}^{\prime}\sim\pi_{\phi}(\cdot|\mathbf{s}^{\prime})}[\bar{Q}_{\theta}(\mathbf{s}^{\prime},\mathbf{a}^{\prime})+{\beta\log\hat{\pi}_{\beta}(\mathbf{a}^{\prime}|\mathbf{s}^{\prime})}]-Q_{\theta}(\mathbf{s},\mathbf{a})\right)^{2}\right].$
Policy:
$\displaystyle~{}~{}\max_{\phi}~{}~{}\mathbb{E}_{\mathbf{s}\sim\mathcal{D},\mathbf{a}\sim\pi_{\phi}(\cdot|\mathbf{s})}\left[Q_{\theta}(\mathbf{s},\mathbf{a})+{\beta\log\hat{\pi}_{\beta}(\mathbf{a}|\mathbf{s})}-{\alpha\log\pi_{\phi}(\mathbf{a}|\mathbf{s})}\right].$
(E.2)

COG. COG (Singh et al., 2020) is an algorithmic framework for utilizing large,
unlabeled datasets of diverse behavior to learn generalizable policies via
offline RL. Similar to real-world scenarios where large unlabeled datasets are
available alongside limited task-specific data, the agent is provided with two
types of datasets. The task-specific dataset consists of behavior relevant for
the task, but the prior dataset can consist of a number of random or scripted
behaviors being executed in the same environment/setting. The goal in this
task is to actually stitch together relevant and overlapping parts of
different trajectories to obtain a good policy that can work from a new
initial condition that was not seen in a trajectory that actually achieved the
reward. COG utilizes CQL as the base offline RL algorithm, and following Singh
et al. (2020), we fix the hyperparameter $\alpha=1.0$ in the CQL part for both
base COG and COG + DR3. All other hyperparameters including network sizes, etc
are kept fixed as the prior work Singh et al. (2020) as well.

####  E.2 Tasks and Environments Used

Atari 2600 games used. For all our experiments, we used the same set of 17
games utilized by Kumar et al. (2021) to test rank collapse. In the case of