# A Moser/Bernstein type theorem in a Lie group with a left invariant metric
under a gradient decay condition

Ari Aiolfi, Leonardo Bonorino, Jaime Ripoll,
Marc Soret, Marina Ville

###### Abstract.

We say that a PDE in a Riemannian manifold $M$ is _geometric_ if, whenever $u$
is a solution of the PDE on a domain $\Omega$ of $M$, the composition
$u_{\phi}:=u\circ\phi$ is also solution on $\phi^{-1}\left(\Omega\right)$, for
any isometry $\phi$ of $M.$ We prove that if $u\in
C^{1}\left(\mathbb{H}^{n}\right)$ is a solution of a geometric PDE satisfying
the comparison principle, where $\mathbb{H}^{n}$ is the hyperbolic space of
constant sectional curvature $-1,$ $n\geq 2,$ and if

$\limsup_{R\rightarrow\infty}\left(e^{R}\sup_{S_{R}}\left\|\nabla
u\right\|\right)=0,$

where $S_{R}$ is a geodesic sphere of $\mathbb{H}^{n}$ centered at fixed point
$o\in\mathbb{H}^{n}$ with radius $R,$ then $u$ is constant. Moreover, given
$C>0,$ there is a bounded non-constant harmonic function $v\in
C^{\infty}\left(\mathbb{H}^{n}\right)$ such that

$\lim_{R\rightarrow\infty}\left(e^{R}\sup_{S_{R}}\left\|\nabla
v\right\|\right)=C.$

The first part of the above result is a consequence of a more general theorem
proved in the paper which asserts that if $G$ is a non compact Lie group with
a left invariant metric, $u\in C^{1}\left(G\right)$ a solution of a left
invariant PDE (that is, if $v$ is a solution of the PDE on a domain $\Omega$
of $G$, the composition $v_{g}:=v\circ L_{g}$ of $v$ with a left translation
$L_{g}:G\rightarrow G,$ $L_{g}\left(h\right)=gh,$ is also solution on
$L_{g}^{-1}\left(\Omega\right)$ for any $g\in G),$ the PDE satisfies the
comparison principle and

$\limsup_{R\rightarrow\infty}\left(\sup_{g\in
B_{R}}\left\|\operatorname*{Ad}\nolimits_{g}\right\|\sup_{S_{R}}\left\|\nabla
u\right\|\right)=0,$

where $\operatorname*{Ad}\nolimits_{g}:\mathfrak{g}\rightarrow\mathfrak{g}$ is
the adjoint map of $G$ and $\mathfrak{g}$ the Lie algebra of $G,$ then $u$ is
constant.

##  1\. Introduction

Finding a characterization of entire minimal graphs on $\mathbb{R}^{n}$ is a
long standing problem on Differential Geometry. The first well known result,
due to Bernstein [2], asserts that the only entire solutions of the minimal
surface equation (MSE) in $\mathbb{R}^{2}$ are affine functions. J. Simons
[12] extended Bernstein’s result to $\mathbb{R}^{n}$ for $2\leq n\leq 7$ and a
celebrated work of Giorgi, Giusti and Bombieri [1] proved that Bernstein’s
theorem is false in $\mathbb{R}^{n}$ for $n\geq 8.$ A full characterization of
entire minimal graphs of $\mathbb{R}^{n}$ was given in terms of the gradient
at infinity of the solutions: Indeed, J. Moser proved that an entire solution
for the MSE in $\mathbb{R}^{n},$ $n\geq 2,$ is an affine function if and only
if the norm of the gradient of the solution is bounded (corollary of Theorem 6
of [7]).

Moser’s result is not true in the hyperbolic space: There are entire solutions
of the MSE in $\mathbb{H}^{n},$ $n\geq 2,$ which assume a prescribed
continuous non-constant value at infinity and with bounded gradient (it is a
consequence, for instance, of Theorem 3.14 of [11]). In this note we prove
that if the norm of the gradient converges exponentially to zero, then the
solution must be constant.

Our main result applies to the solutions of a quite broad class of partial
differential equations, according to the following definitions, and is optimal
in this class. Although we don’t have a proof, we believe that the exponential
decay is also optimal for the MSE.

######  Definition 1.

We say that a partial differential equation (PDE) on a complete Riemannian
manifold $M$ satisfies the _comparison principle_ if, given a bounded domain
$\Omega\subset M,$ if $u$ and $v$ are solutions of the PDE in $\Omega\ $and
$u\leq v$ in $\partial\Omega$ that is

$\limsup_{k}\left(u(x_{k})-v(x_{k})\right)\leq 0$

for any sequence $x_{k}\in\Omega$ that leaves any compact subset of $\Omega,$
then $u\leq v$ in $\Omega.$

######  Definition 2.

We say that a PDE in a Riemannian manifold $M$ is a _geometric_ PDE __ if,
whenever $u$ is a solution of the PDE on a domain $\Omega$ of $M$, the
composition $u_{\phi}:=u\circ\phi$ is also solution on
$\phi^{-1}\left(\Omega\right)$ for any isometry $\phi$ of $M.$

Any PDE depending only on $u,$ $\nabla u,$ $\nabla^{2}u$ and other geometric
operators as $\Delta$ and $\operatorname{div}$ are invariant by isometries and
hence are geometric PDE’s. Those of the form

$\operatorname{div}\left(\frac{a\left(\left\|\nabla
u\right\|\right)}{\left\|\nabla u\right\|}\nabla u\right)+C=0,\text{ }C\text{
constant},$

where $a\in C^{1}\left(\left[0,\infty\right)\right)$, $a^{\prime}>0,$
$a\left(0\right)=0$ also satisfy the comparison principle (Proposition 3.1 of
[11]). This class includes the minimal surface equation with

$a\left(s\right)=\frac{s}{\sqrt{1+s^{2}}}$

and the $p-$Laplace PDE where

$a\left(s\right)=s^{p-1},\text{ }p>1.$

Also the solutions of PDE’s of the form $F\left(\nabla^{2}u\right)=0$ studied
in [3] satisfy the comparison principle (Lemma B of [3]). Comparison
principles for more general geometric PDE’s of the form $F(u,\nabla
u,\nabla^{2}u)=0$ are studied in [9]. On another hand, a PDE depending on
given function on $M$ may not be geometric. For example $\Delta u=f$ is
geometric if and only if $f$ is invariant by the isometries of $M.$ Hence, if
$M$ is a homogeneous manifold then $\Delta u=f$ is geometric if and only if
$f$ is constant. In general, PDE’s depending only on the differentiable
structure of $M$ will not be geometric as $X\left(X\left(u\right)\right)=0,$
where $X$ is a vector field on $M.$ We prove:

######  Theorem 3.

Let $\mathbb{H}^{n}$ be the hyperbolic space of constant sectional curvature
$-1,$ $n\geq 2.$ Let $u\in C^{1}\left(\mathbb{H}^{n}\right)$ be a solution of
a geometric PDE satisfying the comparison principle. If

$\limsup_{R\rightarrow\infty}\left(e^{R}\sup_{S_{R}}\left\|\nabla
u\right\|\right)=0,$

where $S_{R}$ is a geodesic sphere of $\mathbb{H}^{n}$ centered at fixed point
$o\in\mathbb{H}^{n}$ with radius $R,$ then $u$ is constant. Moreover, given
$C>0,$ there is a bounded non constant harmonic function $v\in
C^{\infty}\left(\mathbb{H}^{n}\right)$ such that

$\lim_{R\rightarrow\infty}\left(e^{R}\sup_{S_{R}}\left\|\nabla
v\right\|\right)=C.$

The first part of Theorem 3 is actually a consequence of a more general result
that applies to a Lie group with a left invariant metric and to the solutions
of a broader class of partial differential equations which we call _left
invariant_ PDE’s. To state it we first recall and introduce some notations and
definitions.

Let $G$ be a Lie group, $e$ its neutral element and $\mathfrak{g}=T_{e}G$ its
Lie algebra. Given $g\in G$ denote by $R_{g}$ and $L_{g}$ the right and left
translations of $G,$ $R_{g}(x)=xg$ and $L_{g}(x)=gx,$ $x\in G,$ and by
$C_{g}=L_{g}\circ R_{g}^{-1}$ the conjugation. Let
$\operatorname*{Ad}\nolimits_{g}=d\left(C_{g}\right)_{e}:\mathfrak{g}\rightarrow\mathfrak{g}$
be the adjoint map of $G.$ Given an inner product $\left\langle\text{ ,
}\right\rangle_{e}$ in $\mathfrak{g}$, let $\left\langle\text{ ,
}\right\rangle$ be the left invariant metric of $G$ determined by
$\left\langle\text{ , }\right\rangle_{e},$ namely

$\left\langle u,v\right\rangle_{g}:=\left\langle
d(L_{g}^{-1})_{g}u,d(L_{g}^{-1})_{g}v\right\rangle_{e},\text{ \ }g\in G,\text{
\ }u,v\in T_{g}G.$

######  Definition 4.

We say that a PDE in $G$ is a _left invariant_ PDE __ if, whenever $u$ is a
solution of the PDE on a domain $\Omega$ of $G$, the composition
$u_{g}:=u\circ L_{g}$ is also solution on $L_{g}^{-1}\left(\Omega\right)$ for
any $g\in G.$

Clearly geometric PDE’s on $G$ are left invariant. However, the converse is
not true. For example, considering $\mathbb{R}^{n}$ as a commutative Lie
group, the PDE $\Delta u=X(u),$ where $X$ is a vector field of
$\mathbb{R}^{n}$ is left invariant but not geometric in general. We prove:

######  Theorem 5.

Let $\Omega$ be a bounded $C^{1}$ domain on a Lie group $G$ with a left
invariant metric and let $u\in C^{1}\left(\overline{\Omega}\right)$ be a
solution of a left invariant PDE satisfying the comparison principle on
$\Omega$. Then

(1) $\sup_{\Omega}\left\|\nabla
u\right\|\leq\sup_{g\in\Omega\cup\Omega^{-1}}\left\|\operatorname*{Ad}\nolimits_{g}\right\|\sup_{\partial\Omega}\left\|\nabla
u\right\|,$

where $\Omega^{-1}=\\{g^{-1}\ |\ g\in\Omega\\}$.

Clearly $\left\|\operatorname*{Ad}\nolimits_{g}\right\|=1$ for any $g\in G$ if
the metric of $G$ is bi-invariant. In this case, then, the solution satisfies
the maximum principle for the gradient that is, it holds the equality in (1).

######  Corollary 6.

Let $u\in C^{1}\left(G\right)$ be a solution of a left invariant PDE
satisfying the comparison principle on a non compact Lie group $G$ with a left
invariant metric. If

$\limsup_{R\rightarrow\infty}\left(\sup_{g\in
B_{R}}\left\|\operatorname*{Ad}\nolimits_{g}\right\|\sup_{S_{R}}\left\|\nabla
u\right\|\right)=0,$

where $B_{R}$ is the geodesic ball of $G$ centered at $e$ with radius $R$ and
$S_{R}=\partial B_{R},$ then $u$ is constant.

We prove that if $G$ has negative sectional curvature then the norm of the
adjoint map satisfies the strong maximum principle (Proposition 11).

A well known paper of J. Milnor [6] studies Lie groups with a left invariant
metric. More recently, $3-$dimensional Lie groups with a left invariant metric
have been classified and an explicit description is given in [8]. Finally we
mention that it follows by the Iwasava decomposition that any symmetric space
of noncompact type is isometric to a Lie group with a left invariant metric
[4].

##  2\. Preliminary facts, proofs of Theorem 5 and Corollary 6

We introduce some notation and prove some preliminary results to be used in
Theorem 5 and Corollary 6. Given $h$ and $g$ in $G$, we denote

$R(g)_{h}=\|(R_{g})_{\star,h}\|$

where $(R_{g})_{\star,h}$ is the derivative of $R_{g}$ at the point $h.$

######  Lemma 7.

The quantity $R(g)_{h}$ does not depend on $h$.

So, from now on we denote it

$R(g)=\|(R_{g})_{\star,e}\|$

and introduce, for a subset $S$ of $G$

$R_{S}=\max_{g\in S}R(g).$

###### Proof.

We first notice that the left and right translations commute, i.e., if
$g_{1},g_{2},h\in G$,

$L_{g_{1}}\circ R_{g_{2}}(h)=R_{g_{2}}\circ L_{g_{1}}(h)=g_{1}hg_{2}.$

Let $h,g\in G$. Then

$R(g)_{h}^{2}=\sup_{X\in
T_{h}G,\|X\|=1}\langle(R_{g})_{\star}X,(R_{g})_{\star}X\rangle_{hg}$
$=\sup_{X\in
T_{h}G,\|X\|=1}\langle(L_{(hg)^{-1}})_{\star}(R_{g})_{\star}X,(L_{(hg)^{-1}})_{\star}(R_{g})_{\star}X\rangle_{e}$
$=\sup_{X\in
T_{h}G,\|X\|=1}\langle(L_{g^{-1}})_{\star}(L_{h^{-1}})_{\star}(R_{g})_{\star}X,(L_{g^{-1}})_{\star}(L_{h^{-1}})_{\star}(R_{g})_{\star}X\rangle_{e}$
(2) $=\sup_{X\in
T_{h}G,\|X\|=1}\langle(L_{g^{-1}})_{\star}(R_{g})_{\star}(L_{h^{-1}})_{\star}X,(L_{g^{-1}})_{\star}(R_{g})_{\star}(L_{h^{-1}})_{\star}X\rangle_{e}.$

For $X\in T_{h}G$, $Z=(L_{h^{-1}})_{\star}X$ belongs to $T_{e}G$; moreover

$\|X\|_{h}=\|(L_{h^{-1}})_{\star}X\|_{1}$

and we can rewrite (2) as

$(\ref{avant-derniere})=\sup_{Z\in
T_{e}G,\|Z\|=1}\langle(L_{g^{-1}})_{\star}(R_{g})_{\star}Z,(L_{g^{-1}})_{\star}(R_{g})_{\star}Z\rangle_{e}$
$=\sup_{Z\in
T_{e}G,\|Z\|=1}\langle(R_{g})_{\star}Z,(R_{g})_{\star}Z\rangle_{g}=(R_{g})_{\star,e}.$

∎

######  Lemma 8.

If $g\in G$,

$R(g)=\|\operatorname*{Ad}\nolimits_{g}\|.$

###### Proof.

If $X\in T_{e}G$, by definition of the metric

$\langle(R_{g})_{\star}X,(R_{g})_{\star}X\rangle_{g}=\langle(L_{g^{-1}})_{\star}(R_{g})_{\star}X,(L_{g^{-1}})_{\star}(R_{g})_{\star}X\rangle_{e}$
$=\langle\operatorname*{Ad}\nolimits_{g}(X),\operatorname*{Ad}\nolimits_{g}(X)\rangle_{e}.$

∎

Denote by $d$ the Riemannian distance in $G.$ The next lemma is elementary and
its proof is therefore omitted:

######  Lemma 9.

Let $\Omega$ be a $C^{1}$ open subset of $G.$ Let $u\in
C^{1}\left(\overline{\Omega}\right)$ and $x_{0}\in\overline{\Omega}$ be given
and assume that $\left\|\nabla u\right\|\left(x_{0}\right)\neq 0.$ Let
$\gamma:\left[0,\varepsilon\right)\rightarrow\Omega$ be an arc length geodesic
such that $\gamma(0)=x_{0}$ and

$\gamma^{\prime}(0)=\frac{\nabla u\left(x_{0}\right)}{\left\|\nabla
u\left(x_{0}\right)\right\|}\text{ or }\gamma^{\prime}(0)=-\frac{\nabla
u\left(x_{0}\right)}{\left\|\nabla u\left(x_{0}\right)\right\|}.$

Then

$\left\|\nabla u\right\|\left(x_{0}\right)=\lim_{t\rightarrow
0}\frac{\left|u\left(\gamma(t)\right)-u\left(x_{0}\right)\right|}{d(\gamma(t),x_{0})}.$

######  Lemma 10.

Let $a,b,g$ be elements in $G$. Denoting by $d$ the distance on $G$, we have

$d(a.g,b.g)\leq R(g)d(a,b)$ $d(a,b)\leq R(g^{-1})d(a.g,b.g).$

###### Proof.

If $\gamma$ is a minimizing geodesic between $a$ and $b\ $then
$d(a,b)=\operatorname*{length}(\gamma).$ Also, since $R_{g}\circ\gamma$ is a
path between $a.g$ and $b.g$,

$d(a.g,b.g)\leq\operatorname*{length}(R_{g}\circ\gamma)\leq
R(g)\operatorname*{length}(\gamma)=R(g)d(a,b).$

We have thus proved the first inequality of the lemma. We derive

$d(a,b)=d\left((a.g)g^{-1},(b.g)g^{-1}\right)\leq R(g^{-1})d(a.g,b.g)$

which proves the second inequality. ∎

###### Proof of Theorem 5.

The proof is an extension of Lemma 12.7 of [5] (see also [13]). Let $\Omega$
be a bounded $C^{1}$ domain on $G$ with and let $u\in
C^{1}\left(\overline{\Omega}\right)$ be a solution of a left invariant PDE
satisfying the comparison principle. Take

$k=\sup\left\\{\frac{\left|u(x)-u(z)\right|}{d(x,z)},\text{ }x\in\Omega,\text{
}z\in\partial\Omega\right\\}.$

From Lemma 9 it is then enough to prove that

(3) $\left|u(x)-u(y)\right|\leq kd(x,y)$

for all $x,y\in\Omega.$ Given $x_{1},x_{2}\in\Omega,$ set $z=x_{1}\cdot
x_{2}^{-1},$

$\Omega_{z}=\left\\{z^{-1}\cdot x\in G\text{ $|$
}x\in\Omega\right\\}=\left\\{x\in G\text{ $|$ }z\cdot x\in\Omega\right\\}$

and define $u_{z}\in C^{1}\left(\overline{\Omega}_{z}\right)$ by
$u_{z}(x)=u(zx).$ We have $\Omega\cap\Omega_{z}\neq\varnothing$ since $x_{2}$
belongs to both. By the comparison principle

$\sup\left\\{u(x)-u_{z}(x)\text{ $|$
}x\in\Omega\cap\Omega_{z}\right\\}=\sup\left\\{u(x)-u_{z}(x)\text{ $|$
}x\in\partial\left(\Omega\cap\Omega_{z}\right)\right\\}.$

Then, in particular

$u(x_{2})-u(x_{1})=u(x_{2})-u_{z}(x_{2})\leq\sup\left\\{u(x)-u(zx)\text{ $|$
}x\in\partial\left(\Omega\cap\Omega_{z}\right)\right\\}.$

Let $x_{0}\in\partial\left(\Omega\cap\Omega_{z}\right)$ be such that

$u(x_{0})-u(zx_{0})=\sup\left\\{u(x)-u(zx)\text{ $|$
}x\in\partial\left(\Omega\cap\Omega_{z}\right)\right\\}.$

If $x_{0}\in\partial\left(\Omega\cap\Omega_{z}\right)$ then either
$x_{0}\in\partial\Omega$ or $zx_{0}\in\partial\Omega.$ It then follows from
the hypothesis that

$u(x_{2})-u(x_{1})\leq kd(x_{0},zx_{0}).$

Now, using Lemma 10 twice, we have

$d(x_{0},(x_{1}x_{2}^{-1})x_{0})\leq R(x_{0})d(1,x_{1}x_{2}^{-1})\leq
R(x_{0})R(x_{2}^{-1})d(x_{2},(x_{1}x_{2}^{-1})x_{2})$
$=R(x_{0})R(x_{2}^{-1})d(x_{2},x_{1}).$

∎

###### Proof of Corollary 6.

Given $R>0$, if $g\in S_{R}$ then

$R=d(e,g)=d(g^{-1},e)=d(e,g^{-1})$

so that $S_{R}^{-1}=S_{R}.$ It follows that $B_{R}^{-1}=B_{R},$ where $B_{R}$
is the geodesic ball centered at $e$ with radius $R.$ Then, if $u\in C^{1}(G)$
is an entire solution satisfying the hypothesis of Corollary 6, it follows
from Theorem 5 that

$\max_{B_{R}}\left\|\nabla u\right\|\leq\max_{g\in
B_{R}}\left\|\operatorname*{Ad}\nolimits_{g}\right\|\max_{S_{R}}\left\|\nabla
u\right\|,$

from which the proof follows. ∎

We close this section by proving a strong maximum principle for the adjoint
map:

######  Proposition 11.

Let $G$ be a Lie group with negative sectional curvature and let $\Lambda$ be
any open subset of $G.$ Then

$\left\|\operatorname*{Ad}\nolimits_{h}\right\|<\sup_{g\in\partial\Lambda}\left\|\operatorname*{Ad}\nolimits_{g}\right\|$

for all $h\in\Lambda.$ In particular

$\sup_{g\in\Lambda}\left\|\operatorname*{Ad}\nolimits_{g}\right\|=\sup_{g\in\partial\Lambda}\left\|\operatorname*{Ad}\nolimits_{g}\right\|.$

###### Proof.

By contradiction, assume that

$\left\|\operatorname*{Ad}\nolimits_{h}\right\|=\sup_{g\in\Lambda}\left\|\operatorname*{Ad}\nolimits_{g}\right\|$

for some $h\in\Lambda.$ There exists $x\in T_{e}G,$ $\left\|x\right\|=1,$ such
that
$\left\|\operatorname*{Ad}\nolimits_{h}\right\|=\left\|\operatorname*{Ad}\nolimits_{h}\left(x\right)\right\|.$
Let $X$ be the right invariant vector field of $G$ such that
$X\left(e\right)=x.$ Choose $u\in T_{h}G,$ $\left\|u\right\|=1,$ such that
$u\neq X\left(h\right)$ and let $\gamma\left(t\right),$ $t\geq 0,\ $be the
geodesic parametrized by arc length in $G$ such that $\gamma\left(0\right)=h$
and $\gamma^{\prime}\left(0\right)=u.$ Then

(4)
$\displaystyle\left.\frac{d}{dt}\left\|\operatorname*{Ad}\nolimits_{\gamma\left(t\right)}\left(x\right)\right\|^{2}\right|_{t=0}$
$\displaystyle=0$ (5)
$\displaystyle\left.\frac{d^{2}}{dt^{2}}\left\|\operatorname*{Ad}\nolimits_{\gamma\left(t\right)}\left(x\right)\right\|^{2}\right|_{t=0}$
$\displaystyle\leq 0.$

We have

$\displaystyle\left\|\operatorname*{Ad}\nolimits_{\gamma\left(t\right)}\left(x\right)\right\|^{2}$
$\displaystyle=\left\langle\operatorname*{Ad}\nolimits_{\gamma\left(t\right)}\left(x\right),\operatorname*{Ad}\nolimits_{\gamma\left(t\right)}\left(x\right)\right\rangle$
$\displaystyle=\left\langle
d\left(L_{\gamma\left(t\right)}^{-1}\right)_{\gamma\left(t\right)}\left(d\left(R_{\gamma\left(t\right)}\right)_{e}\right)\left(x\right),d\left(L_{\gamma\left(t\right)}^{-1}\right)_{\gamma\left(t\right)}\left(d\left(R_{\gamma\left(t\right)}\right)_{e}\right)\left(x\right)\right\rangle$
$\displaystyle=\left\langle
d\left(R_{\gamma\left(t\right)}\right)_{e}\left(x\right),d\left(R_{\gamma\left(t\right)}\right)_{e}\left(x\right)\right\rangle,\text{
}t\geq 0.$

Then

$\left\|\operatorname*{Ad}\nolimits_{\gamma\left(t\right)}\left(x\right)\right\|^{2}=\left\langle
X\left(\gamma\left(t\right)\right),X\left(\gamma\left(t\right)\right)\right\rangle=\left\|X\left(\gamma\left(t\right)\right)\right\|^{2},\text{
}t\geq 0.$

Since right invariant vector fields are Killing fields and Killing fields
restricted to geodesics are Jacobi vector fields,

$J\left(t\right):=X\left(\gamma\left(t\right)\right)$

is a Jacobi field along $\gamma,$ $t\geq 0.$ Thus,
$\left\|\operatorname*{Ad}\nolimits_{\gamma\left(t\right)}\left(x\right)\right\|^{2}$
is equal to the square of the norm of the Jacobi field $J\left(t\right)$ along
$\gamma$ satisfying the initial conditions

$\displaystyle J\left(0\right)$ $\displaystyle=X\left(h\right)$ $\displaystyle
J^{\prime}\left(0\right)$
$\displaystyle=\nabla_{\gamma^{\prime}\left(0\right)}X=\nabla_{u}X.$

We have, from (4)

$\left.\frac{d}{dt}\left\|J\left(t\right)\right\|^{2}\right|_{t=0}=2\left\langle
J^{\prime}\left(0\right),J\left(0\right)\right\rangle=0.$

And, using the Jacobi equation,

$\displaystyle\frac{d^{2}}{dt^{2}}\left\|J\left(t\right)\right\|^{2}$
$\displaystyle=2\left\langle
J^{\prime\prime}\left(t\right),J\left(t\right)\right\rangle+2\left\langle
J^{\prime}\left(t\right),J^{\prime}\left(t\right)\right\rangle$
$\displaystyle=-2\left\langle
R(\gamma^{\prime},J)\gamma^{\prime},J\left(t\right)\right\rangle+2\left\|J^{\prime}\left(t\right)\right\|^{2}.$

Since $u\neq X\left(h\right)$ the vector fields $J(t)=X(\gamma(t))$ and
$\gamma^{\prime}(t)$ are linearly independent along $\gamma$ and we then have

$\frac{d^{2}}{dt^{2}}\left\|J\left(t\right)\right\|^{2}=-2K\left(\gamma^{\prime},J\right)\left\|\gamma^{\prime}\wedge
J\right\|^{2}+2\left\|J^{\prime}\right\|^{2}.$

Since $K<0$ it follows that

$\left.\frac{d^{2}}{dt^{2}}\left\|J\left(t\right)\right\|^{2}\right|_{t=0}>0$

contradicting (5). This proves the proposition. ∎

##  3\. The hyperbolic space. Proof of Theorem 3.

We begin by calculating the norm of the adjoint map in the hyperbolic space to
apply Corollary 6. We present an explicit construction of the Lie group
structure of the hyperbolic space, that comes from the Iwasawa decomposition
[4].

In the half-space model

$\mathbb{H}^{n}=\left\\{\left(x_{1},...,x_{n}\right),\text{ $|$
}x_{n}>0\right\\},\text{ }ds^{2}=\frac{\delta_{ij}}{x_{n}^{2}},$

of the hyperbolic $n-$dimensional space, $n\geq 2,$ given $s>0$ and
$t:=\left(t_{1},...,t_{n-1}\right)\in\mathbb{R}^{n-1}$, define

$\displaystyle a_{s},n_{t}$
$\displaystyle:\mathbb{H}^{n}\rightarrow\mathbb{H}^{n},$ $\displaystyle
a_{s}\left(x_{1},...,x_{n}\right)$
$\displaystyle=s\left(x_{1},...,x_{n}\right)$ $\displaystyle
n_{t}\left(x_{1},...,x_{n}\right)$
$\displaystyle=\left(x_{1}+t_{1},...,x_{n-1}+t_{n-1},x_{n}\right).$

Set

$G:=AN=\left\\{a_{s}\circ n_{t}\text{ $|$ }s>0,\text{
}t\in\mathbb{R}^{n-1}\right\\}$

where

$\displaystyle A$ $\displaystyle=\left\\{a_{s},\text{ }s>0\right\\}$
$\displaystyle N$ $\displaystyle=\left\\{n_{t},\text{
}t\in\mathbb{R}^{n-1}\right\\}.$

Given $p\in\mathbb{H}^{n}$ there is one and only one $g_{p}\in
G\subset\operatorname*{Iso}\left(\mathbb{H}^{n}\right)$ such that
$p=g_{p}\left(\left(0,...,0,1\right)\right).$ Indeed: If
$p=\left(x_{1},...,x_{n}\right)\in\mathbb{H}^{n}$ define
$n:=n_{\left(x_{1},...,x_{n-1},0\right)},a:=a_{x_{n}}\in\operatorname*{Iso}\left(\mathbb{H}^{n}\right),$
that is

$\displaystyle n\left(z_{1},...,z_{n}\right)$
$\displaystyle=\left(z_{1}+x_{1},...,z_{n-1}+x_{n-1},z_{n}\right)$
$\displaystyle a$ $\displaystyle=x_{n}\left(z_{1},...,z_{n}\right),\text{
}\left(z_{1},...,z_{n}\right)\in\mathbb{H}^{n}.$

Then, taking

(6) $g_{p}=n\circ a$

we have

$g_{p}\left(0,...,0,1\right)=n\left(a\left(0,...,0,1\right)\right)=n\left(0,...,0,x_{n}\right)=\left(x_{1},...,x_{n}\right)=p.$

Note that $n$ and $a$ are not uniquely determined by $p,$ but $g_{p}$ is.

One may see that with the operation

$p\cdot q:=\left(g_{p}\circ g_{q}\right)\left(\left(0,...,0,1\right)\right)$

$\mathbb{H}^{n}$ is a Lie group (a solvable Lie group indeed. This is a
general fact that holds for symmetric spaces of non compact type ([4], Chapter
VI)). Moreover, given $p\in\mathbb{H}^{n}$ we have, for any
$q\in\mathbb{H}^{n},$

$L_{p}\left(q\right)=p\cdot q=\left(g_{p}\circ
g_{q}\right)\left(\left(0,...,0,1\right)\right)=g_{p}\left(g_{q}\left(\left(0,...,0,1\right)\right)\right)=g_{p}\left(q\right)$

that is $L_{p}=g_{p}.$ Since
$g_{q}\in\operatorname*{Iso}\left(\mathbb{H}^{n}\right)$ it follows that the
left translation $L_{p}$ is an isometry of $\mathbb{H}^{n}$ with respect to
the hyperbolic metric that is, the hyperbolic metric is left invariant with
respect to the Lie group structure of $\mathbb{H}^{n}.$

######  Proposition 12.

Consider $\mathbb{H}^{n}$ as a Lie group and let $e$ be its neutral element.
If $B_{R}$ is the closed geodesic ball of $\mathbb{H}^{n}$ centered at $e$
with radius $R$ then

(7) $\max_{g\in B_{R}}\left\|\operatorname*{Ad}\nolimits_{g}\right\|=\cosh
R+\sinh R.$

###### Proof.

We claim that it is enough to consider the $2-$dimensional case. Indeed: Let
$g\in B_{R}$ and $x\in T_{e}\mathbb{H}^{n}$ be such that

$\left\|\operatorname*{Ad}\nolimits_{g}\left(x\right)\right\|=\max_{h\in
B_{R}}\left\|\operatorname*{Ad}\nolimits_{h}\right\|.$

Let $y\in T_{e}\mathbb{H}^{n}$ be a nonzero vector tangent to the geodesic
from $e$ to $g.$ There exists a totally geodesic hyperbolic plane
$\mathbb{H}^{2}$ of $\mathbb{H}^{n}$ such that $e\in\mathbb{H}^{2}$ and
$x,y\in T_{e}\mathbb{H}^{2}.$ Since $\mathbb{H}^{2}$ is a Lie subgroup of
$\mathbb{H}^{n}$ we have

$\left\|\operatorname*{Ad}\nolimits_{g}\left(x\right)\right\|=\max_{h\in
D_{R}}\left\|\operatorname*{Ad}\nolimits_{h}\right\|$

where $D_{R}=B_{R}\cap\mathbb{H}^{2}$ is the closed geodesic disk centered at
$e$ with radius $R.$ It is clear that the maximum of the norm of the adjoint
map does not depend on hyperbolic plane containing $e.$ This proves our claim.

Considering the half-plane model for $\mathbb{H}^{2}$ with
$e=\left(0,1\right),$ given $p\in\mathbb{H}^{2},$ the conjugation
$C_{p}:\mathbb{H}^{2}\rightarrow\mathbb{H}^{2}$ is given by

(8) $C_{{}_{p}}\left(q\right)=\left(g_{p}\circ g_{q}\circ
g_{p}^{-1}\right)\left(0,1\right)=g_{p}\left(g_{q}\left(g_{p}^{-1}\left(0,1\right)\right)\right).$

Expanding (8) we arrive to

$C_{\left(x,y\right)}\left(z,w\right)=\left(-xw+yz+x,w\right)$

from which we obtain, at a given $X:=\left(a,b\right)\in
T_{\left(0,1\right)}\mathbb{H}^{2},$

$\operatorname*{Ad}\nolimits_{\left(x,y\right)}\left(a,b\right)=d\left(C_{\left(x,y\right)}\right)_{e}\left(a,b\right)=\left(-xb+ya,b\right).$

Since at $\left(0,1\right)$ the hyperbolic metric coincides with the Euclidean
metric and since
$\operatorname*{Ad}\nolimits_{\left(x,y\right)}\left(X\right)\in
T_{\left(0,1\right)}\mathbb{H}^{2}$ we obtain

(9)
$\left\|\operatorname*{Ad}\nolimits_{\left(x,y\right)}\left(X\right)\right\|=\sqrt{\left(-xb+ya\right)^{2}+b^{2}}.$

Since the identity of $\mathbb{R}_{+}^{2}$ is a conformal map between the
Euclidean and hyperbolic geometries, the hyperbolic geodesic circle in the
half plane model of $\mathbb{H}^{2}$ centered at $(0,1)$ with hyperbolic
radius $R$ is also an Euclidean circle. Moreover, because an Euclidean
symmetry with respect to a vertical line is also a hyperbolic isometry, the
Euclidean center of the circle is at some point $\left(0,y_{0}\right)$ of the
vertical line $x=0.$ The Euclidean circle is parametrized by

(10) $\left(0,y_{0}\right)+r\left(\cos\theta,\sin\theta\right),\text{
}\theta\in\left[0,2\pi\right),$

where $r$ is the Euclidean radius. Since
$\gamma\left(t\right)=\left(0,e^{t}\right)$ is an arc length hyperbolic
geodesic such that $\gamma\left(0\right)=\left(0,1\right)$ we have
$t=d_{\mathbb{H}^{2}}\left(\left(0,1\right),\left(0,e^{\pm t}\right)\right).$
In particular

$R=d_{\mathbb{H}^{2}}\left(\left(0,1\right),\left(0,e^{R}\right)\right)=d_{\mathbb{H}^{2}}\left(\left(0,1\right),\left(0,e^{-R}\right)\right).$

It follows that the points $\left(0,e^{R}\right)$ and $\left(0,e^{-R}\right)$
are both in the hyperbolic circle centered at $\left(0,1\right)$ and with
hyperbolic radius $R.$ Then the Euclidean circle must contain both points
$\left(0,e^{R}\right)$ and $\left(0,e^{-R}\right).$ Since these points are in
the same vertical straight line, the Euclidean center of Euclidean circle is

$\frac{\left(0,e^{R}\right)+\left(0,e^{-R}\right)}{2}=\left(0,\cosh R\right).$

And the Euclidean radius $r$ of this hyperbolic circle is just half of the
Euclidean distance between the points $\left(0,e^{R}\right)$ and
$\left(0,e^{-R}\right)$ that is,

$r=\frac{1}{2}d_{\mathbb{R}^{2}}\left(\left(0,e^{R}\right),\left(0,e^{-R}\right)\right)=\frac{e^{R}-e^{-R}}{2}=\sinh
R.$

It follows from (10) that the geodesic disk $D_{R}$ of $\mathbb{H}^{2}$
centered at $\left(0,1\right)$ with radius $R$ is given by

$D_{R}=\left\\{\left(\sinh r\cos\theta,\cosh r+\sinh r\sin\theta\right)\text{
$|$ }\theta\in\mathbb{R},\text{ }0\leq r\leq R\right\\}.$

From (9),

$\left\|\operatorname*{Ad}\nolimits_{\left(x,y\right)}\left(X\right)\right\|^{2}=\left(\cosh
r+\sinh r\sin\theta\right)^{2}\left(\frac{-b\sinh r\cos\theta}{\cosh r+\sinh
r\sin\theta}+a\right)^{2}+b^{2}$

and we may see that the biggest value of right hand side of the last
inequality occurs at $\theta=\pi/2$ and $a=1,$ $b=0.$ From Proposition 11
(which, in the case of the hyperbolic space, can also be directly confirmed
from the expression above) we obtain (7). ∎

###### Proof of Theorem 3.

The first part of the proof is a direct consequence of Corollary 6 and
Proposition 12. For the second part consider a polar coordinate system
$(r,\theta,\varphi_{1},\dots,\varphi_{n-2})$ of $\mathbb{H}^{n}$ centered at a
point $o\in\mathbb{H}^{n}$ that is,
$\Theta=\left(\theta,\varphi_{1},...,\varphi_{n-2}\right)$ are spherical
coordinates of the unit sphere $\mathbb{S}^{n-1}$ centered at the origin of
$T_{o}\mathbb{H}^{n}$ and $p\in\mathbb{H}^{n}\backslash\left\\{o\right\\}$ is
parametrized by

$p=\exp_{o}\left(r\Theta\right),\text{ }r>0,\text{
}\Theta\in\mathbb{S}^{n-1}.$

Let $v$ be the function that depends only on $r$ and $\theta$ given by

$v(r,\theta)=\frac{C(n-1)}{2}\cdot\frac{\displaystyle\int_{0}^{r}(\sinh
s)^{n-1}ds}{(\sinh r)^{n-1}}\cos\theta.$

Remind that in this coordinate system, the metric has the form

$\displaystyle ds^{2}=dr^{2}+\sinh^{2}rd\theta$
$\displaystyle+\sinh^{2}r\sin^{2}\theta
d\varphi_{1}+\sinh^{2}r\sin^{2}\theta\sin^{2}\varphi_{1}d\varphi_{2}$
$\displaystyle+\dots+\sinh^{2}r\sin^{2}\theta\sin^{2}\varphi_{1}\dots\sin^{2}\varphi_{n-3}d\varphi_{n-2}$

and, since $v=v(r,\theta)$, the Laplacian of $v$ can be expressed by

$\Delta v=\frac{\partial^{2}v}{\partial\,r^{2}}+(n-1)\coth r\frac{\partial
v}{\partial\,r}+(n-2)\frac{\cot\theta}{\sinh^{2}r}\cdot\frac{\partial
v}{\partial\,\theta}+\frac{1}{\sinh^{2}r}\cdot\frac{\partial^{2}v}{\partial\,\theta^{2}}.$

Hence, by a straight calculation, we have that $\Delta v=0$. Moreover, since
$v$ depends only on $r$ and $\theta$ and the fields
$E_{1}:=\frac{\partial}{\partial\,r}$ and
$E_{2}:=\frac{\partial}{\partial\,\theta}$ are orthogonal, we have

$\displaystyle\|\nabla
v(r,\theta)\|^{2}=\frac{|E_{1}(v)|^{2}}{\|E_{1}\|^{2}}+\frac{|E_{2}(v)|^{2}}{\|E_{2}\|^{2}}=\left(\frac{|\partial_{r}v|^{2}}{1}+\frac{|\partial_{\theta}v|^{2}}{\sinh^{2}r}\right)$
$\displaystyle=\left(\frac{C(n-1)}{2}\right)^{2}\left(1-(n-1)\displaystyle\cosh
r(\sinh r)^{-n}\int_{0}^{r}(\sinh s)^{n-1}ds\right)^{2}\cos^{2}\theta$
$\displaystyle+\left(\frac{C(n-1)}{2}\right)^{2}\left(\displaystyle(\sinh
r)^{-n}\int_{0}^{r}(\sinh s)^{n-1}ds\right)^{2}\sin^{2}\theta.$

Then,

$\|\nabla v(r,\theta)\|=\frac{C}{2(1+\cosh r)}\quad\text{ for any }\;r\geq
0\;\text{ if }\;n=2$

and

$\|\nabla v(r,\theta)\|\approx Ce^{-r}\sin\theta\quad\text{ as
}\;r\rightarrow+\infty\;\text{ if }\;n>2.$

In both cases

$\sup_{S_{R}}\|\nabla v(R,\theta)\|\approx Ce^{-R}\quad\text{as }\quad
R\rightarrow+\infty.$

Therefore, $v$ is a bounded non constant harmonic function such that

$\lim_{R\rightarrow\infty}\left(e^{R}\sup_{S_{R}}\left\|\nabla
v\right\|\right)=C,$

completing the proof. ∎

## References

  * [1] E. Bombieri, E. De Giorgi, E. Giusti: _Minimal cones and the Bernstein problem,_ Inventiones Mathematicae, Vol 7, 243 – 268, 1969.
  * [2] S. N. Bernstein: _Sur une théorème de géometrie et ses applications aux équations dérivées partielles du type elliptique_ , Comm. Soc. Math. Kharkov, Vol 15, N. 2, 38 - 45, 1915 - 1917.
  * [3] L. Caffarelli, L. Nirenberg, J. Spruck: _The Dirichlet problem for nonlinear second order elliptic equations, III: Functions of the eigenvalues of the Hessian_ , Acta Mathematica, Vol 155, N. 1, 261 301, 1985.
  * [4] S. Helgason: _Differential Geometry, Lie Groups, and Symmetric Spaces_ , Graduate Studies in Mathematics, Vol 34, AMS, 2001.
  * [5] E. Giusti: _Minimal surfaces and functions of bounded variation,_ Monographs of Math, Vol 80, Birkhäuser, 1984
  * [6] J. Milnor: _Curvatures of Left Invariant Metrics on Lie Groups,_ Advances in Mathematics, Vol 21. 1976, 293 - 329.
  * [7] J. Moser: _On Harnack’s Theorem for Elliptic Differential Equations,_ Communications on Pure and Applied Mathematics, Vol XIV, 577- 591,1961.
  * [8] W. Meeks, J. Pérez: _Constant mean curvature surfaces in metric Lie groups_ , Contemporary Mathematics, 2011.
  * [9] P. Pucci, J. Serrin: _The maximum principle,_ Birkhäuser Verlag AG, 2007.
  * [10] H. Rosenberg, F. Schulze, J. Spruck: _The half-space property and entire positive minimal graphs in_ $M\times\mathbb{R},$ Journal of Differential Geometry, Vol 95, N. 2, 2013, 321-336.
  * [11] J. Ripoll, F. Tomi:  _Notes on the Dirichlet problem of a class of second order elliptic partial differential equations on a Riemannian manifold,_ Series “Ensaios de Matemática”, Brazilian Mathematical Soc., Vol 32, 2018.
  * [12] J. Simons: _Minimal varieties in Riemannian manifolds_ , Ann. of Math., Vol 88, 62-105, 1968.
  * [13] G. Williams: _The Dirichlet problem for the minimal surface equation,_ Proceedings of the Centre for Mathematics and its Applications, Instructional Workshop on Analysis and Geometry, Part 1, Volume 34-1, 91 - 110, 1996.

Ari Aiolfi
Universidade Federal de Santa Maria
Brazil
<EMAIL_ADDRESS>
Leonardo Bonorino
Universidade Federal do Rio Grande do Sul
Brazil
<EMAIL_ADDRESS>
Jaime Ripoll
Universidade Federal do Rio Grande do Sul
Brazil
<EMAIL_ADDRESS>
Marc Soret

Université de Tours
France

<EMAIL_ADDRESS>
Marina Ville

Univ Paris Est Creteil, CNRS, LAMA, F-94010 Creteil
France

<EMAIL_ADDRESS>


# Offline Imitation Learning from Multiple Baselines
with Applications to Compiler Optimization

Teodor V. Marinov  Alekh Agarwal  Mircea Trofin

###### Abstract

This work studies a Reinforcement Learning (RL) problem in which we are given
a set of trajectories collected with $K$ baseline policies. Each of these
policies can be quite suboptimal in isolation, and have strong performance in
complementary parts of the state space. The goal is to learn a policy which
performs as well as the best combination of baselines on the entire state
space. We propose a simple imitation learning based algorithm, show a sample
complexity bound on its accuracy and prove that the the algorithm is minimax
optimal by showing a matching lower bound. Further, we apply the algorithm in
the setting of machine learning guided compiler optimization to learn policies
for inlining programs with the objective of creating a small binary. We
demonstrate that we can learn a policy that outperforms an initial policy
learned via standard RL through a few iterations of our approach.

Machine Learning, ICML

##  1 Introduction

When applying Reinforcement Learning (RL) to real-world applications, two key
challenges often prove to be critical blockers to adoption. First is that the
online interaction-then-update loop in conventional RL poses a significant
engineering overhead in most large-scale systems, that are more naturally
designed to take a static Machine Learning (ML) model as a dependency and
update this model only periodically in an offline manner. Second is that RL
algorithms typically begin tabula rasa, that is, they only leverage the
information they glean about the task at hand through these online
interactions. Typical scenarios, where the use of RL is often preceded by
prior attempts using rule-based or supervised ML approaches, come with a
treasure trove of valuable data about desirable and undesirable behaviors,
ignoring which leads to undesirable sample complexity of learning from scratch
for RL. More important, the previously tried decision making policies, even
when individually suboptimal, provide a valuable source of insight into the
plausibly good choices in many scenarios. In this work, we study the question
of leveraging such prior policies and any data collected using them, without
necessarily relying on online policy updates.

Given these shortcomings, offline RL (Ernst et al., 2005; Hester et al., 2018;
Kumar et al., 2020; Cheng et al., 2022), where the agent just learns from a
static dataset collected using some arbitrary policy, as well as hybrid
protocols (Song et al., 2022; Haarnoja et al., 2018; Silver et al., 2014)
interpolating the fully online and offline settings have been proposed in the
literature to take advantage of existing datasets, as well as to ease the
requirement of completely online policy updates. In a different thread of
work, the substantial literature on imitation learning (Pomerleau, 1988; Ross
et al., 2011; Abbeel & Ng, 2004; Ho & Ermon, 2016) aims to leverage any
existing policies that we seek to improve upon, along with the data collected
using them. While imitation learning is studied in both online and offline
settings, the particular scenario of having access to multiple policies of
variable qualities that is of interest here, is only previously studied in an
online setting (Cheng et al., 2020; Barreto et al., 2020).

In this paper, we formalize this question of having access to $K$ baseline
policies $\pi_{1},\ldots,\pi_{K}$, each of which can be quite suboptimal in
isolation, and which we hope are strong in complementary parts of the state
space. We further restrict ourselves to only receiving static datasets $D^{i}$
collected from each policy $\pi_{i}$, and seek to learn a policy which can
combine the strengths of all the baseline policies. We are particularly
interested in challenging settings where the underlying RL problem has a long
horizon, and we only receive sparse trajectory-level feedback at the end of
each trajectory. To motivate this setting, we consider a running example of
optimizing the inlining decisions in a compiler. The horizon of the RL problem
here corresponds to the number of callsites in the program or function being
compiled., which can range from tens to tens of thousands. The reward of a
trajectory is the size of the binary we obtain after compiling the entire
function. The offline setting we study is extremely well motivated here, where
each interaction with the RL environment involves the expensive operation of
compiling and linking the program, and integrating this within the RL loop
engenders a significant engineering overhead.. With this setup, our paper
makes the following contributions:

  * •

A natural behavior cloning algorithm, BC-Max, that combines the multiple
policies by executing each policy in every starting state in our dataset, and
then imitating the trajectory of the policy with the highest reward in that
state. We give an upper bound on the expected regret of the learned policy to
the maximal reward obtained in each starting state by choosing the best
baseline policy for that state.

  * •

We complement our analysis with a lower bound showing that the result is
unimprovable beyond polylogarithmic factors in our setting.

  * •

We apply BC-Max to two different real-world datasets for the task of
optimizing compiler inlining for binary size, and show that we outperform
strong baselines in both the cases. In both cases we start with a single
baseline policy, which is a prior model trained using online RL, which already
has a strong performance on this task. We demonstrate the versatility of BC-
Max by iteratively applying BC-Max on the initial expert, along with all prior
policies trained using previous BC-Max iterations as the next set of
baselines. We show that with a limited amount of interaction with the
environment (to collect trajectories using each successive set of baselines),
we obtain strong policies in a small number of iterations, creating a
promising practical recipe for challenging real-world settings.

##  2 Setting and Related Work

We now define the problem setting formally, and then discuss some lines of
prior work which are relevant to this setting.

###  2.1 Problem setting

#### Contextual MDP setting.

We consider a contextual Markov Decision Process (MDP) with state space
$\mathcal{S}$ and action space $\mathcal{A}$. We denote the initial state
distribution $D_{1}$ and the sampled context (initial state) is $x\sim D$.
Once the context is sampled, the transition kernel $\mathbb{P}_{x}$ is
deterministic, that is $\mathbb{P}_{x}(\cdot|s,a)$ is a point-mass
distribution. The reward function is $r_{x}(s,a)$ and we assume deterministic
rewards. We note that both the transition kernel and reward kernel are
context-dependent. We will omit the context subscript from our notation
whenever it does not introduce ambiguity.

We work in the finite-horizon setting and denote the horizon as $H$. The value
function of a policy $\pi$ is

$\displaystyle
V_{\pi}(x)=\sum_{h=1}^{H}\operatorname{{\mathbb{E}}}_{a_{h}\sim\pi(\cdot|s_{h})}[r_{x}(s_{h},a_{h})],$

where $s_{h}$ is the the state at step $h$ s.t.
$\mathbb{P}_{x}(s_{h}|s_{h-1},a_{h-1})=1$ and $s_{1}\equiv x$. Importantly,
the policy class is such that the action distribution at state $s$ depends
only on $s$ and not on the context, that is $\pi(\cdot|s,x)=\pi(\cdot|s)$.

#### Goal.

We assume that we are given a set of $K$ baselines policies
$\\{\pi_{i}\\}_{k\in[K]}$ together with $n$ trajectories for each policy,
which we denote by $\\{\tau_{i,j}\\}_{i\in[K],j\in[n]}$, where $x_{j}\sim D$
and a trajectory for policy $\pi$ consists of $\\{(s_{h},a_{h})\\}_{h\in[H]}$,
with $a_{h}\sim\pi(\cdot|s_{h})$. For an arbitrary policy $\pi$ and context
$x$ we use $\tau_{\pi}(x)$ to denote the trajectory generated by following
$\pi$ on context $x$. We also assume that we see the total reward for each
trajectory and policy, that is for all $i\in[K],j\in[n]$ we only observe

$\displaystyle
r(\tau_{i,j})=\sum_{(s_{j,h},a_{j,h})\in\tau_{i,j}}r(s_{j,h},a_{j,h}),$

instead of observing a dense reward across all the time steps of the
trajectory. We assume that the rewards are bounded and non-negative, that is
$r(\tau)\in[0,B]$ for all trajectories $\tau$ and some constant $B$.

We have access to a policy class $\Pi$ and seek to find a find a policy in
$\Pi$ which ideally competes with each of the baselines, and is able to
combine their strengths. Let $V_{i}(x^{1})=V_{\pi_{i}}(x^{1})$ denote the
expected cumulative reward of baseline $i$, conditioned on the context $x$.
Then we seek to minimize the regret:

$\text{Reg}(\pi):=\operatorname{{\mathbb{E}}}_{x\sim
D}\left[\max_{i=1}^{K}V_{i}(x)-V_{\pi}(x)\right].$ (1)

That is, we seek to compete with the best of the baselines for each individual
context.

#### Learning setup.

We assume access to the policy class $\Pi$, but do not assume any other
function approximators, such as for modeling value functions. This is partly
due to the fact that the typical training of value functions using Bellman
backups is not feasible in our sparse-reward setting. Furthermore, typical
actor-critic techniques make strong completeness and realizability assumptions
on the value function class, which are not realistic with a restricted notion
of state which we encounter in our motivating problem of compiler
optimization. This necessitates the development of purely policy-based
methods.

###  2.2 Related work

#### Vanilla behavior cloning

Behavior cloning (Widrow, 1964; Pomerleau, 1988) refers to the approach of
learning a policy that matches the mapping from states to actions observed in
the data. This is typically solved as a classification problem for
deterministic policies, or by maximizing the log-likelihood of the chosen
actions in the observed states for stochastic policies. It is unclear how to
apply vanilla behavior cloning in the presence of multiple baselines. We will
present a natural formulation to behavior clone the best baseline policy per
context in the following section.

#### Value-based improvement upon multiple baselines (MAMBA)

Cheng et al. (2020) show how to simultaneously improve upon multiple baseline
policies to compete with the best policy _at each state_ in the MDP, which is
a significantly stronger notion that competing with the best baseline in each
context only. However, this comes with two caveats. Their method requires
value function estimation for the baselines and access to the MDP to execute
trajectories under the learner’s policy and/or baselines. Barreto et al.
(2020) also study a problem which involves improving over multiple baseline
policies, which they title General Policy Improvement. The policy improvement
step again requires value function evaluation. We do not assume such access to
additional function approximators or the MDP in this work.

#### Offline RL:

Without access to the MDP, a natural approach is to consider offline
reinforcement learning, with the data collection policy being a mixture of the
baselines $\pi_{i}$, say chosen uniformly. Given the recent results on offline
RL to compete with any policy that is covered by the data distribution (Kumar
et al., 2020; Xie et al., 2021; Zhan et al., 2022), we can expect a favorable
bound on the regret (1), since all the baselines have a good coverage under
the uniform data collection policy. However, existing offline RL methods with
theoretical guarantees are typically based on function approximation, relying
on actor-critic or $Q$-learning style approaches and on strong credit
assignment using per timestep rewards rather than the aggregated reward of a
trajectory. Applying these techniques using policy-based function
approximation alone and with aggregated reward feedback is not feasible as we
argue through a simple lower bound example.

##  3 Algorithm and Regret Bound

We now describe our algorithm, BC-Max, and give an upper bound on the regret
it incurs to the best per-context baseline.

#### Algorithm.

We describe BC-Max in Algorithm 1. The basic idea of the algorithm is quite
simple. For each context $x_{j}$ in our dataset, we first choose the
trajectory with the highest cumulative reward across all the baselines. Then
we use a standard behavior cloning loss to mimic the choice of actions in this
trajectory. For a context $x_{j},j\in[n]$, we denote
$i_{j}=\operatorname*{\text{argmax}}_{i\in[K]}r(\tau_{i,j})$, and BC-Max tries
to find a policy $\hat{\pi}\in\Pi$ that optimizes the following intuitive
objective:

$\displaystyle\hat{\pi}=\operatorname*{\text{argmin}}_{\pi\in\Pi}\sum_{j=1}^{n}\sum_{(s_{j,h},a_{j,h})\in\tau_{i_{j},j}}\mathbf{1}(\pi(s_{j,h})\neq
a_{j,h}).$ (2)

Algorithm 1 BC-Max for cloning best per-context baseline.

0: Base policies $\\{\pi_{i}\\}_{i\in[K]}$ and policy class $\Pi$.

0: Policy $\hat{\pi}\in\Pi$.

for $j\in[n]$ do

Sample $x_{j}\sim D_{1}$, collect trajectories $\\{\tau_{i,j}\\}_{j\in[n]}$

Compute highest reward policy
$\pi_{i_{j}}=\operatorname*{\text{argmin}}_{i\in[K]}\sum_{(s_{j,h},a_{j,h})\in\tau_{i,j}}r(s_{j,h},a_{j,h})$.

end for

$\hat{\pi}=\operatorname*{\text{argmin}}_{\pi\in\Pi}\sum_{j=1}^{n}\sum_{(s_{j,h},a_{j,h})\in\tau_{i_{j},j}}\mathbf{1}(\pi(s_{j,h})\neq
a_{j,h})$.

One natural question at this point might be that if there are two trajectories
with very similar high rewards in a context, can it help to leverage this
information rather than only picking the one with the higher reward and
cloning it. This is indeed a shortcoming of BC-Max, and other behavior cloning
style approaches. However, we note that we only have access to a trajectory-
level reward, and hedging between two very different trajectories can create a
very noisy learning setup for the algorithm. In situations where value-
functions can be feasibly learned, such information is naturally modeled
through the value function which assigns similar future rewards to similarly
good actions, but we do not find a natural way for incorporating this
information in our setup.

#### Performance guarantee for BC-Max.

We now give a bound on the suboptimality of the policy learned by BC-Max,
relative to the best per-context baseline, in terms of the rewards. The
analysis mirrors the standard results for behavior cloning algorithms (Ross &
Bagnell, 2010). We begin with a realizability assumption which governs how
well the best per-context baseline can be approximated using the learner’s
policy class $\Pi$.

######  Assumption 3.1.

Let
$\tau^{*}(x)=\operatorname*{\text{argmax}}_{\tau_{\pi_{i}}(x)}r(\tau_{\pi_{i}}(x))$
be the trajectory with maximum return over all policies $\pi_{i},i\in[K]$.
There exists $\pi^{*}\in\Pi$ such that

$\displaystyle\mathbb{P}_{x\sim
D}(\tau_{\pi^{*}}(x)\neq\tau^{*}(x))\leq\epsilon.$

Here $\mathbb{P}_{x\sim D}(A)$ denotes the probability of an event $A$ under
the distribution $D$, which we recall is the distribution over the contexts
$x$. The assumption is natural as BC-Max cannot do a good job of approximating
the best per-context baseline when no policy in the policy class has a small
error in achieving this task. Note that the assumption does not take rewards
into account as BC-Max only matches the actions of $\tau^{*}(x)$, and does not
reason about the reward sub-optimality of other actions, as is common in
behavior cloning setups. Indeed this assumption is unavoidable in our problem
setting as we illustrate in the next section.

######  Theorem 3.2.

Under Assumption 3.1, after collecting $n$ trajectories from each of the $K$
base policies, Algorithm 1 returns a policy $\hat{\pi}$ with regret at most

$\displaystyle\text{Reg}(\hat{\pi})\leq O\bigg{(}\epsilon
H+\frac{H^{2}\log(H|\Pi|/\delta)}{n}\bigg{)},$

with probability at least $1-\delta$.

###### Proof.

Recall the definitions of $\tau^{*}(x)$ from Assumption 3.1, and let
$\pi^{*}=\operatorname*{\text{argmin}}_{\pi\in\Pi}\operatorname{{\mathbb{E}}}_{x\sim
D}(\sum_{h=1}^{H}\mathbf{1}(\pi(s(x))\neq a(x)))$ where
$(s_{h}(x),a_{h}(x))_{h=1}^{H}=\tau^{*}(x)$ form the best trajectory for $x$
among the baseline policies. Under Assumption 3.1, we know that
$\operatorname{{\mathbb{E}}}_{x\sim D}(\sum_{h=1}^{H}\mathbf{1}(\pi(s(x))\neq
a(x)))\leq\epsilon H$. Let us define

$\displaystyle\hat{A}(\pi)=$
$\displaystyle\sum_{j=1}^{n}\sum_{h=1}^{H}\mathbf{1}(\pi(s_{j,h})\neq
a_{j,h}),$ $\displaystyle A(\pi)=$
$\displaystyle\operatorname{{\mathbb{E}}}_{x}\left(\sum_{h=1}^{H}\mathbf{1}(\pi(s_{h}(x))\neq
a_{h}(x))\right).$

Clearly we have that $\operatorname{{\mathbb{E}}}[\hat{A}(\pi)]=A(\pi)$ for
any fixed policy $\pi$, and $\hat{A}(\pi)=\sum_{j=1}^{n}Z_{j}(\pi)$ with
$Z_{j}(\pi)=\sum_{h=1}^{H}\mathbf{1}(\pi(s_{j,h})\neq a_{j,h})\geq 0$. We note
that the $Z_{j}$ are i.i.d., with
$\operatorname{{\mathbb{E}}}[Z_{j}(\pi)]=A(\pi)$ and
$\operatorname{{\mathbb{E}}}[Z_{j}(\pi)^{2}]\leq
H\operatorname{{\mathbb{E}}}[Z_{j}(\pi)]=H\,A(\pi)$. Then by Bernstein’s
inequality combined with a union bound over policies, we have with probability
at least $1-\delta$, for all $\pi\in\Pi$:

$\displaystyle|\hat{A}(\pi)-nA(\pi)|\leq$
$\displaystyle\sqrt{nH\,A(\pi)\log(2|\Pi|/\delta)}+H\log(2|\Pi|/\delta)$
$\displaystyle\leq$
$\displaystyle\frac{nA(\pi)}{2}+\frac{3}{2}H\log(2|\Pi|/\delta).$

Applying the inequality with $\pi=\hat{\pi}$ and $\pi=\pi^{*}$, we obtain

$\displaystyle A(\hat{\pi})\leq$
$\displaystyle\frac{2}{n}\hat{A}(\hat{\pi})+\frac{3H\log(2|\Pi|/\delta)}{n}$
$\displaystyle\frac{1}{n}\hat{A}(\pi^{*})\leq$
$\displaystyle\frac{3}{2}A(\pi^{*})+\frac{3}{2}\frac{H\log(2|\Pi|/\delta)}{n}.$

Scaling the second inequality by 2 and adding them yields

$\displaystyle A(\hat{\pi})\leq
3A(\pi^{*})+\frac{6H\log(2|\Pi|/\delta)}{n}\leq
3\epsilon+\frac{6H\log(2|\Pi|/\delta)}{n},$ (3)

where the second inequality follows by Assumption 3.1.

Now we note that for any policy $\pi$:

$\displaystyle\text{Reg}(\pi)$
$\displaystyle=\operatorname{{\mathbb{E}}}_{x}[\max_{i}V_{i}(x)-V_{\pi}(x)]=\operatorname{{\mathbb{E}}}_{x}[r(\tau(x))-r(\tau_{\pi}(x))]$
$\displaystyle\leq\sum_{h=1}^{H}(H-h)\operatorname{{\mathbb{E}}}_{x}[\mathbf{1}(\pi(s_{h}(x))\neq
a_{h}(x))|].$

Plugging the bound from Equation 3 into the inequality above completes the
proof. ∎

#### Implementation details

In practice we can not directly compute $\hat{\pi}$ as defined in Equation 2.
Instead we solve a proxy to the optimization problem by replacing the
indicator function $\mathbf{1}(\pi(s_{j,h})\neq a_{j,h})$ by the cross-entropy
loss. Let $y_{j,h}\in\\{0,1\\}^{A}$ be the indicator with only entry equal to
$1$ the one which corresponds to the action $a_{j,h})$, and all other entries
equal to $0$. Further, we assume that all $\pi\in\Pi$ are such that
$\pi(s)\in\Delta^{A-1}$, that is each $\pi(s)$ represents a distribution over
the actions that policy $\pi$ plays when in state $s$. We then use a first
order method to minimize the loss

$\displaystyle\min_{\pi\in\Pi}\sum_{j=1}^{n}w_{j}\sum_{h=1}^{H}\sum_{a}-y_{j,h}(a)\log(\pi(S_{j,h})),$

where $w_{j}\in[0,1]$ are example weights which we find helpful in our
practical implementation. We refer the reader to our experimental evaluation
for mode details on how the weights are induced.

##  4 Lower bounds

In this section, we show a series of lower bounds which illustrate the
necessity of various aspects of our guarantee in Theorem 3.2. We start with
the necessity of Assumption 3.1

#### Necessity of Assumption 3.1

Let us consider a contextual multi-armed bandit problem, meaning that we fix
$H=1$. For any $\epsilon$, we choose the context space $\mathcal{S}=[M]$ for
$M=\lceil\frac{1}{\epsilon}\rceil$, and choose $D$ to be the uniform
distribution on $\mathcal{S}$. We fix ${\mathcal{A}}=\\{a_{1},a_{2},a_{3}\\}$,
and $K=1$ with the data collection policy $\pi_{1}$ choosing $a=a_{1}$ for
each context $x\in\mathcal{S}$. We consider two possible environments, defined
through rewards $r_{1},r_{2}$. For $a_{1}$, we have
$r_{1}(x,a_{1})=r_{2}(x,a_{2})=1$. For the other two actions, we have
$r_{1}(x,a_{2})=0$ and $r_{1}(x,a_{3})=1$, while the second environment has
$r_{2}(x,a_{2})=1$ and $r_{2}(x,a_{3})=0$. We design a policy class $\pi$ with
two policies $\\{\pi_{1},\pi_{2}\\}$ such that $\pi_{1}(x)=\pi_{2}(x)=a_{1}$
for all $x\neq 1$ and $\pi_{1}(1)=a_{2}$m $\pi_{2}(1)=a_{3}$. Clearly this
policy class satisfies Assumption 3.1. But it also contains an optimal policy
for both the rewards $r_{1},r_{2}$ with a regret equal to $0$. However, since
the data contains no information about which one of $r_{1}$ or $r_{2}$
generated the data, the best we can do is to pick between $\pi_{1}$ and
$\pi_{2}$ uniformly at random, and incur a regret of at least $0.5\epsilon$.
This argument shows that we cannot replace the $0-1$ loss for measuring the
accuracy of a policy in Assupmtion 3.1, with a more reward-aware quantity. It
is also evident from the example that we cannot avoid incurring a regret of
$\Omega(\epsilon)$.

#### Necessity of the comparator choice

As discussed earlier, we define regret relative to the best baseline per-
context, but the broader literature on offline RL allows stronger benchmarks,
such as the best policy covered by the data generating policies. To understand
the difference between the two, we consider the case of $H=2$ and $K=2$, with
$\mathcal{S}=\\{x\\}$ being a singleton. Suppose we have two actions
${\mathcal{A}}=\\{a_{1},a_{2}\\}$ and the baselines choose the trajectories
$\pi_{1}(x)=a_{1}$ and $\pi_{2}(x)=a_{2}$ at both $h=1,2$. There are two
possible reward functions given by
$r_{1}((a_{1},a_{1}))=r_{1}((a_{2},a_{2}))=0.5$, $r_{1}((a_{1},a_{2}))=1$,
$r_{1}((a_{2},a_{1}))=0$ and $r_{2}(\tau)=1-r_{1}(\tau)$ for all trajectories
$\tau$. Now we observe that in the sense of the coverage studied in the
offline RL literature, all four trajectories are covered by the dataset, since
we get to observe both the actions at both the steps in the episode. However,
since the data contains no useful information to distinguish between $r_{1}$
and $r_{2}$, no learning method can pick a covered policy which is better than
our benchmark of best baseline per-context. This challenge arises in our
scenario as we only observe the aggregate reward over a trajectory, which
makes per-step credit assignment used in standard offline RL methods through
the use of Bellman errors infeasible.

#### Tightness of horizon factor

The tightness of the horizon factor follows from a simple reduction to Theorem
6.1 in Rajaraman et al. (2020), who show that in a finite horizon episodic
MDP, there is no algorithm which only observes $n$ optimal policy trajectories
and returns a policy with regret better than
$\Omega\left(\frac{SH^{2}}{n}\right)$. The MDP contructed by Rajaraman et al.
(2020) for the lower bound has a non-deterministic transition kernel
$\mathbb{P}$. Since we are assume that the transitions are deterministic, we
instead use the randomness in sampling contexts, $x\sim D$ to simulate
$\mathbb{P}$. Concretely, let $\xi_{1},\ldots,\xi_{H}$ be the random bits
sampled at the $H$ steps in a fixed episode from the construction in Rajaraman
et al. (2020). We define $x=(\xi_{1},\ldots,\xi_{H})$ and set
$\mathbb{P}_{x}(s_{h+1}|s_{h},a_{h})=\mathbb{P}(s_{h+1}|s_{h},a_{h},\xi_{h})$
to be the trasition taken for the realization of $\xi_{h}$. The size of the
state space of contextual MDP with the above transition kernel can be set to
$S=\Theta(\log_{2}(|\Pi|)$ and the action space to
${\mathcal{A}}=\\{a_{1},a_{2}\\}$, so that each policy in $\Pi$ is encoded by
how it acts on the state space. This construction directly leads to the
following lower bound.

######  Theorem 4.1.

For any number of samples $n$ there exists a family of contextual MDPs with
disribution over contexts given by $D_{1}$, such that the policy
$\hat{\pi}=\texttt{A}(\\{\tau_{\pi^{*}}(x_{i})\\}_{i=1}^{n})$ returned by any
algorithm A satisfies

$\displaystyle\operatorname{{\mathbb{E}}}_{x}[V_{\pi^{*}}(x)]-\operatorname{{\mathbb{E}}}_{x}[V_{\hat{\pi}}(x)]\geq\min\left\\{H,\frac{\log_{2}(|\Pi|)H^{2}}{n}\right\\}.$

##  5 Case study: Optimizing a compiler’s inlining policy

###  5.1 The inlining for size problem

In short the inlining problem which we study in our experiments consists of
deciding to inline or not to inline a callsite in a program with the goal of
minimizing the size of the final program binary. We omit most compilation
details and just give a brief overview which should be sufficient for
understanding the problem from a RL perspective. In our specific scenario,
compilation is split into a frontend (fe) and a backend (be). In our setup the
frontend consists of translating the program into an Intermediate
Representation (IR), doing some frontend optimizations, then a (thin) link
step follows, which re-organizes functions in the various modules to improve
inlining opportunities. For more details on the linking step see Johnson et
al. (2017). The backend compilation follows the thin link step and is applied
on the updated modules. It consists of further optimizations, including
inlining decisions, final linking and lowering the IR to machine code, e.g.,
x86, ARM, etc. The IRs with which our RL algorithms work with are post
frontend linking and pre backend optimization, that is we work only on backend
optimizations. We note that a program is made up of multiple _modules_. In the
fe, a module corresponds to a single C/C++ source file, after all the
preprocessor directives have been applied. In the be, the module would consist
of a mix of IRs from different fe modules. The inlining decisions will be
taken at callsites in the IRs of each module, where the callee is also in the
same module. We note that each module is ultimately compiled to a machine
code-specific binary, with its own size, that will still need to be linked
into the final executable. Hence we treat each module as a context $x$ in our
contextual MDP setting, and the value $V_{i}(x)$ of the baseline policy
$\pi_{i}$ is the size of the binary we get for module $x$, when we make
inlining decisions for the module according to $\pi_{i}$.

$\displaystyle\texttt{Program}\to\texttt{IRs}\to\texttt{{fe optimizations}}$
(4) $\displaystyle\to$
$\displaystyle\texttt{ThinLinking}\xrightarrow{\texttt{{Collecting
IRs}}}\texttt{{be optimization}}$ $\displaystyle\to$
$\displaystyle\texttt{Final linking}\to\texttt{x86}$

In Equation 4 we outline the compilation process, together with the step at
which we collect IRs from the respective modules to be used in our RL
algorithms. It is important to note that the learned RL policy will make
inlining decisions both at the fe optimization and be optimization parts in
Equation 4, however, the IRs for training are only collected after the fe
optimization step.

The contextual MDP setting can now be tied together with the compilation
process as follows. Each context is a module as mentioned above, with state-
space defined by its IR. Each state corresponds to a callsite in the IR of the
module $x$, and the action set is $\\{\texttt{inline},\texttt{ don't
inline}\\}$ or $\\{1,0\\}$ respectively. Each $\pi_{i}$ is some base inlining
policy and $V_{i}(x)$ is defined by the size of the compiled stand-alone
module $x$. It is important to note that there is a mismatch between trying to
maximize $\operatorname{{\mathbb{E}}}_{x\sim D}[V_{\pi}(x)]$ and the overall
goal of minimizing the binary size, as it is not necessarily true that the sum
of module sizes equals the size of the binary. In fact, part of the post-
inlining be and linker optimizations may introduce a significant distribution
shift between the sum of module sizes and the size of the final binary. In our
experiments, we try to minimize this distribution shift by turning off certain
optimizations. For more details on the compilation pipeline we refer to Trofin
et al. (2021).

We note that the entire process is fully deterministic, as we assumed in our
theoretical setup, since the compiler is a deterministic program.

###  5.2 Dataset collection

We train and evaluate on two sets of binaries. In the first experiment we
train on a proprietary search binary and evaluate the model on a different
proprietary set of targets that are part of a cloud systems infrastructure.
These targets need to be installed on a fixed size partition of each cloud
machine and hence are size-constrained. In the second experiment we train and
evaluate on the Chrome binary on Android. Training proceeds in two separate
steps, which we repeat over several iterations. The two steps can be
summarized as follows, first we collect a training dataset which consists of
trajectories with smallest size over all base policies available at the
current iteration. Next, we train a new base model using the objective defined
in Equation 2. This conceptually applies Algorithm 1 repeatedly, where the set
of baseline policies is updated at each iteration to include the new policy
obtained from the previous iteration. We now describe each step carefully.

#### Training dataset collection.

The dataset collection begins by creating a corpus of IRs of modules which
make up the final binary. The corpus creation follows the work of Trofin et
al. (2021); Grossman et al. (2023) and uses tools for extracting the corpus
are available on GitHub222A detailed example can be found at
https://github.com/google/ml-compiler-opt/blob/main/docs/inlining-
demo/demo.md. The corpus is created at the beginning of training and remains
the same throughout every iteration. Training begins under the assumption that
there exists at least one base policy. In the first iteration a training
dataset is collected from this initial base policy $\pi_{1}$, next, $\pi_{1}$
is behavior cloned by solving the optimization problem in Equation 2. We
specify how the weights for the objective are computed in the following
sections as they are different for the different targets. Let $\hat{\pi}_{2}$
denote the resulting policy after solving Equation 2. This policy is non-
deterministic and so we construct the base policy $\pi_{2}$ by setting
$\pi_{2}(s)=\operatorname*{\text{argmax}}_{a\in\\{0,1\\}}\hat{\pi}_{2}(s,a)$,
that $\pi_{2}$ always plays the most likely action according to
$\hat{\pi}_{2}$. This concludes the first iteration. More generally, if we
have a larger initial set of baseline policies than just a
singleton$\\{\pi_{1}\\}$, the iterations proceed similarly. However, instead
of just using $\pi_{1}$, we use the full set $\\{\pi_{i}\\}_{i\in[K]}$ of
baseline policies at every iteration at the first iteration.

Proceeding this way, at the $t$-th iteration the set of base policies is taken
as a subset of $\\{\pi_{1},\ldots,\pi_{t-1}\\}$ which always contains
$\pi_{1}$ (or the larger set of all initial baselines). Then we again invoke
BC-Max with these baseline policies, and obtain a new randomized policy
$\hat{\pi}_{t}$, and we refer to $\pi_{t}$ as the corresponding deterministic
greedy policy. When collecting a new training dataset we not only collect
trajectories with the chosen subset of base policies but we also may force
exploration by using $\hat{\pi}_{t-1}$ in the way discussed next.

#### Exploration in training dataset collection.

For a fixed module $x$ and a policy $\pi$, we choose a ceiling on the number
of exploration steps as a hyper parameter, which is a fraction of the length
of the trajectory $|\tau_{\pi_{1}}(x)|$. The call-sites at which exploration
occurs are selected as follows. The first exploration call-site is selected as
$\tilde{h}=\operatorname*{\text{argmin}}_{h}\\{|\hat{\pi}(s_{h})(0)-\hat{\pi}(s_{h})(1)|\\}_{S_{h}\in\tau_{\pi}(x)}$,
as the call-site where the exploration policy $\hat{\pi}$ is the least
confident about the action to choose. The exploration step is then played at
$s_{\tilde{h}}$ by taking the action $1-\pi(s_{\tilde{h}})$ (recall that
$\pi(s)\in\\{0,1\\}$), and the remaining steps in the trajectory are completed
by playing according to $\pi$. Let $\hat{\tau}$ denote the trajectory from the
last round of exploration. In the following exploration round the exploration
step is selected as the step $h$ at which the gap,
$|\hat{\pi}(s_{h})(0)-\hat{\pi}(s_{h})(0)|$, is smallest among all
$h>\tilde{h}$, where $\tilde{h}$ is the exploration step in the previous
round. Once the maximum number of exploration rounds is reached or the
exploration step reaches the end of the trajectory, we return the trajectory
which results in the smallest module size among all explored trajectories.
Pseudo-code is presented in Algorithm 2. The exploration strategy is governed
by $\hat{\pi}_{t-1}$, however, it can be updated by using the non-
deterministic policy $\hat{\pi}$ which induces $\pi$. We leave such approaches
as future work, as we have already observed significant benefit of using only
$\hat{\pi}_{t-1}$ as the exploration policy.

Algorithm 2 Explore module

0: Base policy $\pi$, exploration policy $\hat{\pi}$, module $x$, maximum
exploration steps $T$.

0: Compilation trajectory $\tau_{\pi}(x)$ with reward $r_{\pi,x}$.

Compute vanilla trajectory $\tau_{\pi}(x)$ by compiling with $\pi$ and receive
reward $r^{1}_{\pi,x}$

$t=1$

$\hat{\tau}_{1}=\tau_{\pi}(x)$

$\tilde{h}_{1}=\operatorname*{\text{argmin}}_{h}\\{|\hat{\pi}(S_{h})(0)-\hat{\pi}(S_{h})(1)|\\}_{S_{h}\in\hat{\tau}_{1}}$

while $t\leq T$ do

Replay $\hat{\tau}_{t}$ until $\tilde{h}_{t}$

Play $1-\pi(S_{\tilde{h}_{t}})$ at $\tilde{h}_{t}$

Complete trajectory $\hat{\tau}^{t+1}$ by playing $\pi$

Receive reward $r^{t+1}_{\pi,x}$

if $\tilde{h}_{t}<|\hat{\tau}_{t+1}|$ then

$\tilde{h}_{t+1}=\operatorname*{\text{argmin}}_{h>\tilde{h}_{t}}\\{|\hat{\pi}(S_{h})(0)-\hat{\pi}(S_{h})(1)|\\}_{S_{h}\in\hat{\tau}^{t+1}}$

else

break

end if

end while

$t^{*}=\operatorname*{\text{argmax}}_{t}r^{t}_{\pi,x}$,
$r_{\pi,x}=r^{t^{*}}_{\pi,x}$, $\tau_{\pi}(x)=\hat{\tau}_{t^{*}}$

#### Online versus offline learning.

Our theoretical setup frames the problem in an offline learning scenario, yet
Algorithm 2 and the iterative procedure do rely on our ability to interact
with the environment in an adaptive manner. Note, however, that the modality
of interaction used in our approach is quite different, and significantly more
practical than full-fledged online RL. Each round of policy learning, which
happens using BC-Max, is fully offline. This process, which involves a large
number ($10^{5}-10^{6}$) stochastic gradient steps, happens without any
interaction with the environment, and is where the bulk of the learning
happens. Subsequently, we form a new data collection policy for the next
iteration, and this policy is applied to collect one trajectory per module.
The data collection process does not involve any policy updates, and hence is
massively parallelizable with no interlocking bottlenecks with the learning
process. In online RL, on the other hand, data collection and policy updating
go hand-in-hand, which typically requires significantly more complex
architecture (Mnih et al., 2016) to scale to domains where data collection is
expensive. Our approach, on the other hand, simply requires interleaving
standard supervised learning and batch data collection, which is quite
desirable especially in the compiler application, where the ML training
happens on GPUs, while the compilation happens on CPU machines.

###  5.3 Search application targets

Similarly to Trofin et al. (2021) we collect a corpus for training purposes
from a search application binary with approximately $30000$ modules. The
initial base policy is an RL model trained using an Evolutionary Strategy
(ES333The policy can be found here: https://github.com/google/ml-compiler-
opt/releases/tag/inlining-Oz-v1.1) as in Trofin et al. (2021). After
collecting a training dataset with the ES policy we noticed that the
distribution of sizes of modules is fairly non-uniform, with few modules
having very large sizes or very small sizes and majority of modules being
somewhere in-between. Because we expect that the actions of the behavior
cloning policy taken on larger size modules are more important for size saving
we upweight the actions in such trajectories. The weights used in training are
computed as follows. Let $\texttt{size}(x,\pi_{1})$ denote the size of module
$x$ from the collected trajectory under policy $\pi_{1}$ (or in the case of
multiple baseline policy under the best baseline policy). The modules are
partitioned into buckets according to their sizes where the limits of the
buckets are taken to be on exponentially scaling grid, that is the first
bucket contains all modules with size $\texttt{size}(x,\pi_{1})\in[0,2^{0})$,
the second bucket all modules such that
$\texttt{size}(x,\pi_{1})\in[2^{0},2^{1})$ etc., up to the final bucket with
size $[2^{M-1},2^{M})$. Let
$b_{m}=\\{x:\texttt{size}(x,\pi_{1})\in[2^{m-1},m)\\}$ denote the $m$-th
bucket and let $m(x)$ be the $m$ for which $x\in b_{m(x)}$. The weight $w_{x}$
for module $x$ is computed as follows.

$\displaystyle w_{x}=\frac{\max_{m}|b_{m}|}{|b_{m(x)}|}.$

Algorithm 3 BC-Max for cloning best per-context baseline with exploration

0: Base policy $\pi_{1}$ and policy class $\Pi$. Max exploration steps $T$.
Max number of iterations $N$.

0: Policy $\hat{\pi}\in\Pi$.

$l=1$

while $l\leq N$ do

for $j\in[n]$ do

Sample $x_{j}\sim D_{1}$

for $i\in\\{\pi_{s}\\}_{s\leq l}$ do

$r_{i,j},\tau_{i,j}=\texttt{Algorithm~{}\ref{alg:explore_module}}(\pi_{i},\hat{\pi}_{\ell},x_{j},T)$

end for

Compute highest reward policy
$\pi_{i_{j}}=\operatorname*{\text{argmin}}_{i\in[K]}\sum_{(s_{j,h},a_{j,h})\in\tau_{i,j}}r_{i,j}(s_{j,h},a_{j,h})$.

end for

Compute module weights $\\{w_{j}\\}_{j\in[n]}$ (See Sect. 5.3,5.4).

$\hat{\pi}_{l}=\operatorname*{\text{argmin}}_{\pi\in\Pi}-w_{j}\sum_{j=1}^{n}\sum_{(s_{j,h},a_{j,h})\in\tau_{i_{j},j}}a_{j,h}\log(\pi(a_{j,h}|s_{j,h}))$

$\pi_{l}(s)=\operatorname*{\text{argmax}}_{a}\hat{\pi}_{l}(a|s),\forall s\in
S$

end while

We train two sets of policies, one set is trained without exploration and is
precisely in line with Algorithm 1. For full pseudo-code, which includes the
exploration step, we refer the reader to Algorithm 3. The second set is
trained with exploration as described in Section 5.2.

Figure 1: Savings in MB from ES on training binary

In Figure 1 we show savings of the trained policies to $\pi_{1}$, which is the
ES policy, on the search binary from which the training dataset is collected.
In Figure 2 we show the savings on a _different test binary_. On the $x$-axis
of the figures we show the size savings of the policy $\pi_{i}$ learned at
each iteration $i$, with and without exploration respectively, where $bc_{0}$
is the behavior cloned policy from ES. Both figures demonstrate the success of
our approach in improving significantly beyond the initial baseline, as well
as the benefits from multiple iterations of the process. Furthermore, the gap
between the lines with and without exploration highlights the benefits of the
added exploration.

We note that the compilation for both the training and test binaries is
carried out in the following way to minimize the distribution shift – the fe
optimizations are carried out by ES, while the be optimizations are carried
out by the trained policies. If we were to use the trained policies in both fe
and be, this might lead to significant distribution shift, as Algorithm 1
works only on trajectories collected after the fe optimizations for which ES
is always used. That is, if any of the trained policies, $bc_{i}$, act very
differently on the fe, compared to ES, the resulting IRs before the be
optimization might be completely different from the training set IRs, and
hence the trained policy might take very sub-optimal actions.

Figure 2: Savings in MB from ES on test binary

###  5.4 Chrome on Android

In our second set of experiments we train an RL policy for Chrome on Android.
The training and test binaries are the same in this case. The base policy with
which we start is an RL policy trained using Proximal Policy Optimization
(PPO444The policy can be found here:
https://commondatastorage.googleapis.com/chromium-browser-
clang/tools/mlgo_model2.tgz) (Schulman et al., 2017) as done in Trofin et al.
(2021). There are two differences in training from Section 5.3. First, we
focus only on the setting where we do exploration. The second difference in
training is how the weights for the objective in Equation 2 are formed. The
approach for computing the weights used here is inspired by the fact that we
want to improve on PPO in each module and not just on the sum of module sizes.
That is we want to maximize the size savings over the worst case module in our
dataset. The following approach is natural when such max-min guarantees are
desired.

Reusing notation from Section 5.3 we let

$\displaystyle p_{x}^{1}$ $\displaystyle=\frac{|b_{m(x)}|}{\sum_{m}|b_{m}|}$
$\displaystyle w_{x}^{1}$
$\displaystyle=\frac{\max_{m}p_{x}^{1}}{p_{m(x)}^{1}},$

be the weights in the first iteration of training. In following iterations the
weights are set as $w_{x}^{t}=\frac{\max_{m}p_{x}^{t}}{p_{m(x)}^{t}}$, where
$p^{t}$ is update using the Hedge algorithm (Littlestone & Warmuth, 1994). The
update uses the sum of sizes in each bucket as losses, normalized by the
$\ell$-infinity norm, that is

$\displaystyle\tilde{L}^{t}_{m}$ $\displaystyle=\sum_{x\in
b_{m}}\texttt{size}(x,\pi_{t})$ $\displaystyle L^{t}_{m}$
$\displaystyle=\frac{\tilde{L}^{t}_{m}}{\|L^{t}\|_{\infty}},$

where $L_{m}^{t}$ denotes the $m$-th coordinate of the loss vector $L^{t}$.
The Hedge update is then

$\displaystyle\tilde{p}^{t+1}_{m}$ $\displaystyle=p^{t}_{m}\exp(-\eta
L^{t}_{m})$ $\displaystyle p^{t+1}_{m}$
$\displaystyle=\frac{\tilde{p}^{t}_{m}}{\sum_{m}\tilde{p}^{t}_{m}}.$

Figure 3: Savings in MBs from PPO on sum of module sizes Figure 4: Savings in
MBs from PPO on binary size

In Figures 3 and 4, we plot the savings of our learned policies across
iterations, relative to the initial PPO policy, measured in two different
ways. For Fig 3, we simply add up the sizes of the binaries produced by
compiling each module in our training dataset. This is a clean metric, as the
distribution shift between training and evaluation is small, and no artifacts
from linker or post-inlining be optimizations are introduced in the
evaluation. As we see, we improve rapidly beyond the PPO policy with the
iterative applications of BC-Max. Note that even the sum of module sizes
suffers from the typical distribution shift between online and offline RL,
since the data used from behavior cloning is collected using a different
policy than the one we apply in evaluation. For the sum of module sizes
metric, we can study the effect of this distribution shift rather carefully by
also compiling with an _oracle_ policy, which simply chooses the best baseline
policy for each module, which is the target for training in BC-Max. This
oracle, shown in red in Figure 3 naturally provides larger gains relative to
PPO than our learned policy as expected, but the gap reduces through the
iterations of our process, indicating that the policies tend to stabilize
through iterations, and the training data for later applications of BC-Max is
closer to on-policy data. We note that the oracle changes between different
instantiations of our weights. This is because the $i$-th trained policy
$\pi_{i}$ depends on the choice of weights and so the oracle after the $i$-th
iteration which chooses the best among $\\{\pi_{i}\\}_{\ell=1}^{i}$ also
depends on the choice of weights. We also note that the gap between the
learned and oracle policy’s performance is smaller when we use the Hedge
weights, and that the weighted version has a bigger gain relative to PPO,
showing the efficacy of this approach.

Finally, in Figure 4 we present the savings in size of the Chrome on Android
binary, which is the actual yardstick. Here we cannot easily evaluate the size
of the oracle, so we only compare our policies to PPO, and again observe
impressive gains, with the Hedge-weighted variant doing better. The binary
size when compiled with the PPO policy is approximately $213.32$ MB.

## Acknowledgements

We would like to thank Ziteng Sun for great discussions on the lower bound.

## References

  * Abbeel & Ng (2004) Abbeel, P. and Ng, A. Y.  Apprenticeship learning via inverse reinforcement learning.  In _Proceedings of the twenty-first international conference on Machine learning_ , pp. 1, 2004.
  * Barreto et al. (2020) Barreto, A., Hou, S., Borsa, D., Silver, D., and Precup, D.  Fast reinforcement learning with generalized policy updates.  _Proceedings of the National Academy of Sciences_ , 117(48):30079–30087, 2020.
  * Cheng et al. (2020) Cheng, C.-A., Kolobov, A., and Agarwal, A.  Policy improvement via imitation of multiple oracles.  _Advances in Neural Information Processing Systems_ , 33:5587–5598, 2020.
  * Cheng et al. (2022) Cheng, C.-A., Xie, T., Jiang, N., and Agarwal, A.  Adversarially trained actor critic for offline reinforcement learning.  In _International Conference on Machine Learning_ , pp. 3852–3878. PMLR, 2022.
  * Ernst et al. (2005) Ernst, D., Geurts, P., and Wehenkel, L.  Tree-based batch mode reinforcement learning.  _Journal of Machine Learning Research_ , 6, 2005.
  * Grossman et al. (2023) Grossman, A., Paehler, L., Parasyris, K., Ben-Nun, T., Hegna, J., Moses, W., Diaz, J. M. M., Trofin, M., and Doerfert, J.  Compile: A large ir dataset from production sources.  _arXiv preprint arXiv:2309.15432_ , 2023.
  * Haarnoja et al. (2018) Haarnoja, T., Zhou, A., Hartikainen, K., Tucker, G., Ha, S., Tan, J., Kumar, V., Zhu, H., Gupta, A., Abbeel, P., et al.  Soft actor-critic algorithms and applications.  _arXiv preprint arXiv:1812.05905_ , 2018.
  * Hester et al. (2018) Hester, T., Vecerik, M., Pietquin, O., Lanctot, M., Schaul, T., Piot, B., Horgan, D., Quan, J., Sendonaris, A., Osband, I., et al.  Deep q-learning from demonstrations.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 32, 2018.
  * Ho & Ermon (2016) Ho, J. and Ermon, S.  Generative adversarial imitation learning.  _Advances in neural information processing systems_ , 29, 2016.
  * Johnson et al. (2017) Johnson, T., Amini, M., and Li, X. D.  Thinlto: scalable and incremental lto.  In _2017 IEEE/ACM International Symposium on Code Generation and Optimization (CGO)_ , pp. 111–121. IEEE, 2017.
  * Kumar et al. (2020) Kumar, A., Zhou, A., Tucker, G., and Levine, S.  Conservative q-learning for offline reinforcement learning.  _Advances in Neural Information Processing Systems_ , 33:1179–1191, 2020.
  * Littlestone & Warmuth (1994) Littlestone, N. and Warmuth, M. K.  The weighted majority algorithm.  _Information and computation_ , 108(2):212–261, 1994.
  * Mnih et al. (2016) Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T., Harley, T., Silver, D., and Kavukcuoglu, K.  Asynchronous methods for deep reinforcement learning.  In _International conference on machine learning_ , pp. 1928–1937. PMLR, 2016.
  * Pomerleau (1988) Pomerleau, D. A.  Alvinn: An autonomous land vehicle in a neural network.  _Advances in neural information processing systems_ , 1, 1988.
  * Rajaraman et al. (2020) Rajaraman, N., Yang, L., Jiao, J., and Ramchandran, K.  Toward the fundamental limits of imitation learning.  _Advances in Neural Information Processing Systems_ , 33:2914–2924, 2020.
  * Ross & Bagnell (2010) Ross, S. and Bagnell, D.  Efficient reductions for imitation learning.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pp. 661–668. JMLR Workshop and Conference Proceedings, 2010.
  * Ross et al. (2011) Ross, S., Gordon, G., and Bagnell, D.  A reduction of imitation learning and structured prediction to no-regret online learning.  In _Proceedings of the fourteenth international conference on artificial intelligence and statistics_ , pp. 627–635. JMLR Workshop and Conference Proceedings, 2011.
  * Schulman et al. (2017) Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ , 2017.
  * Silver et al. (2014) Silver, D., Lever, G., Heess, N., Degris, T., Wierstra, D., and Riedmiller, M.  Deterministic policy gradient algorithms.  In _International conference on machine learning_ , pp. 387–395. Pmlr, 2014.
  * Song et al. (2022) Song, Y., Zhou, Y., Sekhari, A., Bagnell, J. A., Krishnamurthy, A., and Sun, W.  Hybrid rl: Using both offline and online data can make rl efficient.  _arXiv preprint arXiv:2210.06718_ , 2022.
  * Trofin et al. (2021) Trofin, M., Qian, Y., Brevdo, E., Lin, Z., Choromanski, K., and Li, D.  Mlgo: a machine learning guided compiler optimizations framework.  _arXiv preprint arXiv:2101.04808_ , 2021.
  * Widrow (1964) Widrow, B.  Pattern recognition and adaptive control.  _IEEE Transactions on Applications and Industry_ , 83(74):269–277, 1964.
  * Xie et al. (2021) Xie, T., Cheng, C.-A., Jiang, N., Mineiro, P., and Agarwal, A.  Bellman-consistent pessimism for offline reinforcement learning.  _Advances in neural information processing systems_ , 34:6683–6694, 2021.
  * Zhan et al. (2022) Zhan, W., Huang, B., Huang, A., Jiang, N., and Lee, J.  Offline reinforcement learning with realizability and single-policy concentrability.  In _Conference on Learning Theory_ , pp. 2730–2775. PMLR, 2022\.