Aligning Agent Policy with Principal Interests:
Reward Design via Bilevel RL

Souradip Chakraborty 1 Amrit Singh Bedi 1 Alec Koppel 2 Dinesh Manocha 1 Furong Huang 1 Mengdi Wang 3

1. Introduction

Abstract

The increasing complexity and widespread use of artificial
agents highlight the critical need to ensure that theirperformance aligns with broader objectives such as social
welfare and economic impacts, which are referred to byeconomists as externalities (Rahwan, 2018). This has
given rise to the field of AI alignment (Liu et al., 2022),
which focuses on aligning agent behavior with societal
goals and values. We study this problem in the context
of reinforcement learning (RL) (Sutton & Barto, 1998),
which may be mathematically posed as a Markov Decision
Problem (Puterman, 2014). This framework gained traction
for its ability to solve games (Mnih et al., 2013; Silver et al.,2017), personalized recommendation systems (Steck et al.,2021), various continuous control tasks (Ng et al., 1999;
Lillicrap et al., 2015), and develop strategies in financial
markets (Ardon et al., 2021). The question of how to alignan RL agent to broader goals may be decomposed into two
major questions:(i) How can one align the behavior of learning agents withbroader objectives?
(ii) Does alignment compromise the performance oflearning agents?

In reinforcement learning (RL), a reward function
is often assumed at the outset of a policy optimization procedure. Learning in such a fixed reward paradigm in RL can neglect important policy
optimization considerations, such as state spacecoverage and safety. Moreover, it can fail to encompass broader impacts in terms of social welfare, sustainability, or market stability, potentiallyleading to undesirable emergent behavior and potentially misaligned policy. To mathematicallyencapsulate the problem of aligning RL policy
optimization with such externalities, we consider
a bilevel optimization problem and connect it toa principal-agent framework, where the principal
specifies the broader goals and constraints of thesystem at the upper level and the agent solves
a Markov Decision Process (MDP) at the lower
level. The upper-level deals with learning a suitable reward parametrization corresponding to thebroader goals and the lower-level deals with learning the policy for the agent. We propose Principal
driven Policy Alignment via Bilevel RL (PPABRL), which efficiently aligns the policy of theagent with the principal’s goals. We explicitly
analyzed the dependence of the principal’s trajectory on the lower-level policy, prove the convergence of PPA-BRL to the stationary point of the
problem. We illuminate the merits of this framework in view of alignment with several examples
spanning energy-efficient manipulation tasks, social welfare-based tax design, and cost-effective
robotic navigation.

To address (i), a litany of approaches have been studied:inverse reinforcement learning (Ziebart et al., 2008; Brown
et al., 2019; Arora & Doshi, 2020) and imitation learning (Ho & Ermon, 2016; Kang et al., 2018; Ghasemipour
et al., 2019; Xiao et al., 2019) employ expert demonstrations to learn a policy or reward function (Torabi et al.,2018; Ouyang et al., 2022). Offline approaches leveragelarge-scale prior-collected data to learn behaviors deemed
appropriate by an expert (Yang et al., 2023; Levine et al.,2020; Yin et al., 2022; Chen et al., 2022). However, acquiring a substantial amount of human feedback or well-aligned
demonstrations can be expensive or infeasible (Bai et al.,2022; Chen et al., 2023a; Wolf et al., 2023), especiallywhen dealing with diverse (Jabbari et al., 2017), personalized, or competing goals set by individuals, organizations,or governments (Bai et al., 2022; Chen et al., 2023a; Liuet al., 2022; Sun et al., 2023). Therefore, instead, we adopt
an approach where one seeks to directly design the data-*

Equal contribution 1 Department of Computer Science, University of Maryland, College Park, USA. 2 JP Morgan Chase AI
Research, USA. 3 Department of Electrical Engineering, Princeton
University/Deepmind, Princeton, NJ, USA. Correspondence to:Souradip Chakraborty <schakra3@umd.edu>.
Interactive Learning with Implicit Human Feedback Workshop at
ICML 2023, Honolulu, Hawaii, USA. PMLR 202, 2023. Copyright2023 by the author(s).

1Submission and Formatting Instructions for ICML 2023

generating mechanism according to some external utility
function. Doing so introduces a dependency on the question(ii), as alignment upends the standard training mechanism
for learning agents. In particular, evaluating an alignment
objective is contingent upon the agent’s performance over atime horizon through trajectories generated by the agent’s
policy. This entanglement naturally leads to a variant of
bilevel optimization.

sociated with the transition dynamics conditioned on fixed
policy at the inner-stage. We develop an algorithmic framework that addresses these issues and therefore enables one todesign RL agents that are aligned with externalities. Therefore, our main contributions are to:• formulate the agent policy alignment problem as a
bilevel optimization problem where the outer objective
centers on reward design through policy evaluation
over the horizon, and the inner level pertains to policyalignment with the designed reward via policy optimization in an MDP;

To date, problems of this type have mostly been studied
through the lens of bilevel optimization (Bracken & McGill,
1973). Bilevel optimization considers a setup where theouter-level objective is a function of decision variables andthe optimizer at the inner-level in both deterministic (Hansen
et al., 1992; Shi et al., 2005) or a probabilistic sense (Chenet al., 2021). However, these works primarily do not allow
either stage to be defined by an MDP.

• by analyzing local KKT points using differentiationand the Implicit function theorem, we derive an iterative procedure to jointly solve for the design parameters at the outer level and policy parameters at the
inner level in this framework. This procedure, known
as Principal-driven Policy Alignment via Bilevel RL(PPA-BRL);

Distinct lines of effort exist from game theory and economics in terms of mechanism design (Myerson, 1989; Hurwicz, 2003; Maskin, 2008), Stackleberg games (Von Stackelberg, 2010), and contract theory (Green & Stokey, 1983;
Eisenhardt, 1989; Stiglitz, 1989) that have a bearing on our
technical approach. Closer to this work is recent intersections of MDPs and algorithmic Stackleberg solvers, which
impose special information structures such as linear MDPs
(Zhong et al., 2021) or oracle access to transition models(Goktas et al., 2022). Most similar to our work are thosethat develop implicit function-theorem based gradient play
(Fiez et al., 2020; Vu et al.); however, these techniques areexperimentally-based. Worth mentioning is principal-agent
problem in contract theory, which admits a Stacklberg formulation (Zhu et al., 2022; Chen et al., 2023b) in terms of
repeatedly interacting bandits.

• establish the proposed methodology converges to a local KKT point of the problem. The main challenge
lies in deriving and dealing with the expectation of
the gradient of the outer objective, which depends on
the policy of the inner level. This results in an interesting notion of score function with respect to outerparameters crucial for design and alignment.

2. Policy Alignment as Bilevel Formulation
Standard policy optimization. Let us start by considering the Markov Decision Process (MDP) tuple M :={S, A, γ, P, r} (Puterman, 2014), which is a tuple consisting of a state space S, action space A, transition dynamics
P, discount factor γ ∈ (0, 1), and reward r : S × A → R.
Starting from a given state s ∈ S, an agent selects an actiona and transitions to another s′ ∼ P(· | s, a). We hypothesize that the agent follows a stochastic stationary policy thatmaps states to distributions over actions πθ : S → P(A),which is parameterized by a parameter vector θ ∈ Rd . We
can write the standard finite horizon policy optimization
problem as
"H−1
#
X
h
max Vs (θ) := Eγ r(sh , ah ) |, s0 = s , (1)A more extensive discussion of these research threads isavailable in Appendix 8. The main takeaway is that no formulation exists in which the inner-stage is a parameterized
policy optimization problem in an MDP, and the outer-stage
is an expected-value objective over the distribution underlying trajectories defined by an RL policy. This setting,to our knowledge, is most appropriate for formalizing thealignment of RL agents to externalities (Rahwan, 2018) andemergent behavior (Teo et al., 2013).
This generality introduces technical impediments not foundin either prior RL methods or bilevel programming: (i) the
outer objective is parameterized by the inner-stage’s optimal policy, which must be numerically approximated; (ii)to solve the aforementioned numerical approximation, thelocal Karush-Kuhn-Tucker (KKT) conditions of the outerstage exhibit dependence on Jacobians and Hessians thatare interrelated with variables at the inner-stage (Hong et al.,2020; Khanduri et al., 2021; Li et al., 2022); and (iii) the
sampling distribution for the expectation at the outer-stage
depends on trajectories generated from product measure as-θ

h=0

where the expectation is with respect to the stochasticity inthe policy πθ and the transition dynamics P. In (1), we used
notation Vs (θ) for value function in state s to emphasize thedependence on parameters θ. We note that the formulationin (1) utilizes a reward function r (fixed a priori), and welearn a policy corresponding to reward function r. But asdetailed in the introduction, policy learning with a fixed
reward does not allow one to tether the training process2Submission and Formatting Instructions for ICML 2023

to an external objective such as (Ni & Paul, 2019), social
welfare (Balcan et al., 2014), or market stability (Buehler
et al., 2019).

encompasses the goal of aligning the reward function with
objectives such as energy usage (Ni & Paul, 2019), social
welfare (Balcan et al., 2014), market stability (Buehler et al.,2019), or other external factors not under the direct control
of an RL agent. The parameters ν under the designer’s control define the incentive structure for the behavior of an RL
agent, and in that way, exhibits parallel with mechanism
design (Lyu et al., 2022a) and contract theory (Zhu et al.,2022).

Proposed formulation. To enable this capability, we formulate a bilevel optimization problem where the inner levelis a Markov Decision Process (MDP) and The outer levelis an expected utility maximization, whose distribution depends upon the transition model and policy associated withtrajectories at the inner level:maxν

G(ν, θ∗ (ν))∗

s.t. θ (ν) := arg max E

"H −1
ℓ
X

θ

Outer objective. To be specific, we consider a utility at theouter level of the form(2)#
γ h rν (sh , ah ) |, s0 = s ,G(ν, θ∗ (ν)) = D(πθ∗ (ν) ) + Z(ν),

which is comprised of two terms: a quantifier U (πθ∗ (ν) ) of
the merit of design parameters ν ∈ Rn , whereas the second
term Z(ν) represents, e.g., a regularizer or prior on the
distribution over trajectories P(τ ; θ∗ (ν)). More specifically,
the explicit mathematical form of U (πθ∗ (ν) ) decides thequality of policy πθ∗ (ν) by collecting trajectories (denotedby τ ), and via associating a designer’s reward U (τ ), foreach trajectory under policy πθ∗ (ν) given by

where ah ∼ πθ (ah |sh ) and the inner-level goal is to maximize the Hℓ horizon γ-discounted cumulative return of
rewards rν (s, a), i.e., the value function.
Inner Objective. As previously mentioned, the interrelationship of these levels is intended to formalize that adesigner, called a principal in contract theory (Zhu et al.,2022; Zhao et al., 2023) , specifies the incentive structureassociated with the reward function of an RL agent at the
inner-level. To be more precise, note that the reward rν (s, a)in (2) is additionally a parametric function of the designer’s
parameters ν ∈ Rn , in contrast to the standard MDP setting(1). Hence, we can also write the inner-level optimization
problem in (2) as
max Vs (ν, θ) := E
θ

(4)

h=0

"H −1
ℓ
X

D(πθ∗ (ν) ) = EP (τ ;θ∗ (ν)) [U (τ )] =X

U (τ ) · P(τ ; θ∗ (ν)),τ(5)
where P(τ ; θ∗ (ν)) denotes the probability of trajectory τ .
Subsequently, we employ the shorthand notation Eπθ∗ (ν) [·]to denote the expectation
with respect to the distribution unQH
uderlying trajectory τ : h=0
P(sh+1 | sh , ah )πθ∗ (ν) (·|sh ),where Hu is the length of upper level trajectory collected at
the lower level optimal policy. The explicit form of such an
objective, i.e., reward specifications for the designer U , isprovided in detailed examples in the following Subsection
2.1.

#
γ h rν (sh , ah ) |, s0 = s , (3)h=0where ah ∼ πθ (ah |sh ) and the expectation is with respect
to the stochasticity in the policy πθ and the transition dynamics P. In contrast to value function in (1), the value functionin (3) is an implicit function of the designer’s parametersν which controls the reward function rν . For the example
of robotic navigation, rν (s, a) is a function of the distanceto a goal location. Further, to make the policy parameter θ
dependence on the reward parameter ν clear, we write θ(ν)
subsequently. In this work, we restrict focus to the case thatthe optimizer θ∗ (ν) at the inner-level is unique, which mandates that one parameterize the policy in a tabular (Agarwal
et al., 2020; Bhandari & Russo, 2021) or softmax fashion(Mei et al., 2020a); otherwise, at most one can hope for witha policy gradient iteration is to obtain an approximate local
extrema (Zhang et al., 2020). We defer a more technicaldiscussion of this aspect to Section 4. For the upper-levelobjective in (2), we aim to maximize an objective function
G(ν, θ∗ (ν)), which depends on the design parameters ν anda cumulative return of the RL agent evaluated at a given policy, which is an implicit function of policy parameters θ∗ (ν).
The outer objective in (2) addresses the limitations of (1)As we will see in the next section, developing an iterative solver to (2) exhibits some unique technical challenges
not found in prior art on bilevel programming (Ghadimi &Wang, 2018a; Akhtar et al., 2022). Before shifting focus todo so, we present some representative examples of (2).
2.1. Motivating Examples
Example 1: Energy efficient and sustainable design for
robotic manipulation. Consider a robotic manipulation
task where the objective of the agent is to learn an optimal policy to transport components from a fixed position
to a target position ν := (x, y). On the other hand, the designer’s objective is to select the work-bench position ν tominimize the energy consumption of the robotic arm during
the transportation task. Hence, it naturally boils down to the3Submission and Formatting Instructions for ICML 2023

yi . Thus, the bilevel formulation is given by
"H −1
#
u
X
max EP (τ ;θ∗ (ν))
vν (sh , ah )following bilevel problem as"H#uX−ah wh | ah ∼ πθ∗ (ν) (·|sh )max EP (τ ;θ∗ (ν))
ν:=(x,y)

ν

h=1

∗

s.t. θ (ν) := arg max E
θ

"H −1
ℓ
X

(6)
#

s.t. θ∗ (ν) := arg max E

γ rν (sh , ah ) |, s0 = s ,#γ h uν (sh , ah )|, s0 = s .

h=0where ah ∼ πθ (ah |sh ), γ is the discount rate, and θ∗ (ν)
represents the optimal inner policy of the household agent,
which maximizes its expected cumulative return over a time
horizon Hℓ .

h=0where ah ∼ πθ (ah |sh ), Eπθ∗ (ν) denotes the expectation
with respect to the trajectories collected by the lower-level
optimal policy πθ∗ (ν) . In (6), action ah denotes the angularacceleration of the robotic arm, the state is represented bysh = (αh , wh ), αt is the discretized angle, and wh angularvelocity of the robotic arm. we define the transitions as(αt+1 , wt+1 ) = (αh + wh , wh + ah ). The reward of the inner objective rν (sh , ah ) = −λ1 ∥sh − ν∥2 − λ2 ∥wh ∥2 , i.e.,
reward increases as the arm moves closer to the workbenchwith a controlled angular velocity. The outer objective focuses on minimizing the energy emission during transport
and is thus entangled with the trajectories generated under
the optimal policy obtained via the lower-level.

3. Algorithmic Policy Alignment with
Externalities via Bilevel RL
To solve the bilevel problem in (2), we seek to develop
stochastic optimization algorithm similar to (Ghadimi &Wang, 2018a) which requires to evaluate the gradients of
the upper and lower level objective. To do so, we begin by
deriving the gradient of the outer objective ∇ν G[ν, θ∗ (ν)]
from (4) with respect to design parameter ν given byX∇ν G[ν, θ∗ (ν)] = ∇νU (τ ) · P(τ ; θ∗ (ν)) + ∇ν Z(ν)

Example 2: Social welfare aligned tax design for households. Consider the problem of tax design for individual
households while remaining attuned to social welfare, motivated by (Hill et al., 2021). From the household’s perspective, each household seeks to maximize its own utility
uh based on the number of working hours, consumptionof goods, and net worth. Let us denote the accumulated
asset as state sh . At each time step h, the household agent
selects an action ah = (nh , ci,h ), ah ∼ πθ (ah |sh ), where
nh is the number of hours worked, and ci,h is the consumption of good i at a pre-tax price of pi , and θ denotes
the policy parameter. We denote f (sh ) as the reward forthe accumulative asset sh , updated
PM at each time step by
sh+1 = sh + (1 − x)wnh − i=1 ci,h and ν = (x, yi )is the income tax rate and consumption tax rate for goodi. Here we note that x is a uniform tax across all households, whereas yi is a household-specific tax rate. Then thehousehold agent’s utility at time steph is given by the equaQM  ci,h νition uh = f (sh ) − γn2h + i=1 pi (1+y, where thei)product term corresponds to Cobb-Douglas function (Rothet al., 2016). In contrast, the objective of the regulatory
body or government (upper-level) is to maximize the social
welfare vt by adjusting the tax rates ν based on the optimal
policy of the household agent (lower level). Hence, theouter objective representing the social
is defined as QwelfareM
PM ci,hi=1 ci,t yi
Q
vh = g(sh ) + i=1 1+yi + ψ ln
+
wx
h
,
n
M
(1+y )
i=1

"H −1
ℓ
X

θ

h

(7)

h=0

τ

= Eτ [U (τ ) · ∇ν log(P(τ ; θ∗ (ν)))] + ∇ν Z(ν)
HX
u −1
h
i
= EP (τ ;θ∗ (ν)) U (τ ) ·∇ν log πθ∗ (ν) (ah |sh )h=0(8)+ ∇ν Z(ν),
where we have used the log-trick and standard rule of expectation to get the final expression in (8), similar to thestandard derivation of the policy gradient formula (Agarwal et al., 2020), and the score function method in general(Williams, 1992; Sutton et al., 1999). We emphasize two aspects: (a) the score function term ∇ν log πθ∗ (ν) (a|s) in (8),which denotes the gradient of the optimal policy with respect
to the design parameter ν; and (b) the expectation is with respect to the trajectory distribution P(τ ; θ∗ (ν)) generated under policy at the inner-level πθ∗ (ν) given by P(τ ; θ∗ (ν)) :=QHℓ −1ρ(s0 ) h=0πθ∗ (ν) (ah |sh )P(s′h | πθ∗ (ν) (sh , ah ).
True to our knowledge, this is the first time this coupledscore function is appearing in the RL training, which captures the change of optimal policy with respect to the reward
function’s design parameters. The term ∇ν log πθ∗ (ν) (a|s)
is crucial for our setting, as the designer (principal agent,such as a regulatory body or central planner) at the upperlevel can directly control the policy learning by modifying the reward parameters. However, the estimation of∇ν log πθ∗ (ν) (ah |sh ) is nontrivial as it depends on the solution of the lower-level problem in (2), and therefore requires the evaluation of hypergradient ∇ν θ∗ (ν). To seeiwhere g(·) is the reward for the accumulative asset, ψ is apositive constant, and w is the wage rate. The household
agent follows a policy that maximizes its discounted cumulative reward, while the social planner aims to maximize thediscounted total social welfare by tuning the tax rates x and4Submission and Formatting Instructions for ICML 2023

that, let us employ the shorthand notation fh (θ∗ (ν)) :=log πθ∗ (ν) (ah |sh ), we can rewrite the gradient as∇ν fh (θ∗ (ν)) = ∇ν θ∗ (ν)T ∇θ fh (θ∗ (ν)).

Algorithm 1 Principal driven Policy Alignment via Bilevel
RL (PPA-BRL)
1: Input: Reward parametrization ν0 policy initialization(9)

θ0 , upper and lower-level step sizes αu > 0, αℓ > 0
respectively
2: for all t = 0, 1, 2, ..., T − 1 do
3:
for all k = 0, 2, ..., K − 1 do
4:
Sample N trajectories τ ∼ P (τ ; θK (νt )) and estimate policy gradient ∇θ Vs (νt , θK (νt )) from equation (27)5:Update the policy gradient parameter as :From the first order optimality condition for the lower levelobjective, it holds that∇θ Vs (ν, θ∗ (ν)) = 0,(10)which is the gradient of lower-level objective with respect
to parameter θ evaluated at the optimal θ∗ (ν). Now, differentiating again with respect to ν on both sides of (10), we
obtainθk+1 (νt ) = θk (νt ) − αℓ ∇θ Vs (νt , θk (νt ))∇2ν,θ Vs (ν, θ∗ (ν)) + ∇ν θ∗ (ν)∇2θ Vs (ν, θ∗ (ν)) = 0. (11)6:
7:

The above expression would imply that we can write thefinal expression for the gradient in (9) as

end forUpdate the reward parameterization in the upperlevel from equation (26) as :e ν G(νt , θK (νt ))νt+1 ← νt − αu ∇∇ν fh (θ∗ (ν)) = −∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θ∗ (ν))−1(12)

8: end for
9: Output: νT , θ K (νT )∇θ fh (θ∗ (ν)).

We substitute (12) into (8) to write the final expression for
optimal objective for convex G). Once, we have learned
the outer objective in (2):the optimal reward parameter νT , we can utilize that to
"
HX
u −1
obtain optimal policy θ∗ (νT ) aligned with the upper-level∇ν G[ν, θ∗ (ν)] =EP (τ ;θ∗ (ν)) U (τ ) ·[−∇2v,θ Vs (ν, θ∗ (ν)) objectives. But as mentioned earlier, there are impediments
h=0

to implementing Algorithm 2 in practice. Therefore, wedevelop a stochastic algorithm without requiring access toθ∗ (ν) oracle next.
Remark 3.1. We have only presented the analytical formsof the first and second-order information required to obtaina numerical solver for problem in (2). However, in practice, these update directions are unavailable due to theirdependence on distributions P(τ ; θ∗ (ν)) and MDP transition model P. Therefore, only sampled estimates of theexpressions in (26)-(29) are available.

(13)#· ∇2θ Vs (ν, θ∗ (ν))−1 ∇θ fh (θ∗ (ν))]+ ∇ν Z(ν).
From the gradient expression in (13), we note that there
are three intertwined technical challenges for solving the
reward alignment problem efficiently:(i) requirement of access to θ∗ (ν)

4. Convergence Analysis(ii) evaluating Jacobians and Hessians of the lower-levelobjectives

In this section, we analyze the convergence behavior of
Algorithm 1. Since the outer objective G(ν, θ∗ (ν)) is nonconvex with respect to ν, we consider ∇ν G(ν, θ∗ (ν)) as
our convergence criteria and show its convergence to a firstorder stationary point, as well as the convergence of θK (ν)
to θ∗ (ν). Taken together, these constitute a local KKTpoint (Boyd & Vandenberghe, 2004)[Ch. 5] Without loss ofgenerality, our convergence analysis is for the minimization(upper and lower level objectives). We proceed then byintroducing some technical conditions required for our main
results.
Assumption 1 (Lipschitz gradient of outer objective). For
any ν ∈ Rn , the gradient of the outer objective is Lipschitz(iii) unbiasedly sampling trajectories from P(τ ; θ∗ (ν)) =QHℓ −1ρ(s0 ) h=0πθ∗ (ν) (ah |sh )P(s′h | πθ∗ (ν) (sh , ah ) dependent on optimal policy πθ∗ (ν)
To provide intuition, let us temporarily hypothesize that an
oracle provides us access to θ∗ (ν) for given ν. In this case,
we can write a first-order gradient iteration to solve the joint
policy learning and alignment problem (2), which is given
as Algorithm 2. We note that from the standard analysis ofgradient descent, we know that the update in Algorithm 2
would converge to stationary point (for non-convex G and5

Submission and Formatting Instructions for ICML 2023

continuous w.r.t to second argument with parameter Lg , i.e.,
we may write∥∇ν G(ν, θ∗ (ν)) − ∇ν G(ν, θK (ν))∥ ≤ Lg ∥θ∗ (ν) − θK (ν)∥.
(14)(ii) The Hessian of the value function is Lipschitz withparameter as

Assumption 2. For all s ∈ S and a ∈ A, reward function isbounded as r(s, a) ≤ R for all (s, a) and Lipschitz w.r.t toν, i.e., |rν1 (s, a)−rν2 (s, a)| ≤ Lr ∥ν1 −ν2 ∥; , and bounded∀(s, a) i.e rν (s, a) ≤ R.
Assumption 3. The policy πθ is Lipschitz with respect to parameter θ, which implies ∥πθ1 (·|s) − πθ2 (·|s)∥ ≤ Lπ ∥θ1 −θ2 ∥ for all θ1 ̸= θ2 . The score function ∇θ log πθ (a|s) isbounded ∥∇θ log πθ (a|s)∥ ≤ B and Lipschitz:∥∇2θ Vs (ν, θ∗ (ν)) − ∇2θ Vs (ν, θK (ν))∥
′

∗

(18)

K

≤ L ∥θ (ν) − θ (ν)∥,
where, L′ = Lf1 χ1 H2ℓ L2 + Lf1 and Lf1 = L2 Hℓ2 R.
Here, Hℓ is the horizon length at the lower level, χ1 is
a constant defined in (64), R is the maximum reward(cf. Assumption 2), and L2 is the Lipschitz parameterof the Hessian of the policy defined in Assumption 3.

∥∇θ log πθ1 (·|s) − ∇θ log πθ2 (·|s)∥ ≤ L1 ∥θ1 − θ2 ∥(15)(iii) The second order mixed jacobian term
∇2ν,θ Vs (ν, θK (ν)) is Lipschitz continuous w.r.t θi.e∀θ1 ̸= θ2 . Further, the policy parameterization inducesscore function whose Hessian is Lipschitz:∥∇2ν,θ Vs (ν, θ∗ (νt )) − ∇2ν,θ Vs (νt , θK (ν))∥∥∇2θ log πθ1 (·|s) − ∇2θ log πθ2 (·|s)∥ ≤ L2 ∥θ1 − θ2 ∥
(16)

′′

∗

(19)

K

≤ L ∥θ (ν) − θ (ν)∥
where, L′′ = Lf3 χ2 H2ℓ L2 + Lf3 and Lf3 = Lr L1 Hℓ2 .
Here, χ2 is a constant defined in equation (72), andother constants are as defined in statement (ii).

∀θ1 ̸= θ2 .
Assumption 4. Vs (ν, θ) satisfies the Polyak-Lojasiewicz(PL) condition with respect to θ with parameter µ. We
denote {λ(∇2θ Vs (ν, θ))j }dj=1 as the eigenvalues of Hessian matrix ∇2θ Vs (ν, θ). Although, Vs (ν, θ) is non-convex
in θ, but follows the restriction on the eigenvalues as
λ(∇2θ Vs (ν, θ)) ∈ [−ˆl, −µ̂] ∪ [µ̂, ˆl].

The proof of Lemma 4.2 is in Appendix 11.2. To prove thisresult, we start by considering the value function expression, and evaluating it’s gradient, Hessian, and Jacobians.
After expanding each of them, we separate the reward andpolicy-related terms and then utilize the aforementioned
assumptions to upper bound them, respectively.
Lemma 4.3. Let us define the update direction associatedwith the gradient in (13) as

Next, we introduce a few key technical lemmas regarding
the iterates generated by Algorithm 1. We begin by quantifying the distributional drift associated with the transition
model under θK (νt ) as compared with θ∗ (νt ) at the lower
level, which results in a transient effect at the outer level.

ϕ1 (τ ) :=U (τ ) ·Lemma 4.1. Under Assumptions 1 - 4, for trajectory τ =u{sh , ah }H
h=1 , it holds that

Hu
X[−∇2v,θ Vs (ν, θ∗ (ν))

(20)

t=0

· ∇2θ Vs (ν, θ∗ (ν))−1 ∇θ fh (θ∗ (ν))],Hu L2 ∗
Df (P (τ ; θ (ν)), P (τ ; θ (ν))) ≤∥θ (ν) − θK (ν)∥,
2
(17)
∗

Hℓ is the horizon length for the lower level [cf. (2)],
Lr is the reward Lipschitz parameter [cf. Assumption2], and B is the score function bound [cf. Assumption
3].

Kand ϕ2 (τ ) :=U (τ ) ·Hu
X[−∇2v,θ Vs (ν, θK (ν))

(21)

t=0

· ∇2θ Vs (ν, θK (ν))−1 ∇θ fh (θK (ν))].

where Df is the f-divergence between distributions and Lπ
is the Lipschitz parameter (cf. Assum. 3).

Under Assumptions 1 - 4, it holds that

The proof of Lemma 4.1 in provided in Appendix 11.1.
Next, we establish some error bound conditions on key
second-order terms that appear in equations (26)-(29) when
we substitute θ∗ (νt ) by θK (νt ).

∥Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ ) − ϕ2 (τ )]∥ ≤ Hu2 ũγ1 ∥θ∗ (ν) − θK (ν)∥,
(22)
L

L′

′′

where γ1 := κL1 + ν,θ
+ Llπ . Here, ũ is the upper bound2
lπof utility U (τ ) defined in (5), lπ , L1 are policy-related Lipschitz parameters (cf. Assumption 3), L′ and L′′ as definedin Lemma 4.2, and κ mixed condition number defined inequation (92) and Lν,θ upper-bound of the norm of second
order mixed jacobian term, defined in equation (90)Lemma 4.2 (Value function related upper bounds). Under
Assumptions 1 - 4, it holds that(i) The second order Jacobian term ∇2ν,θ Vs (νt , θK (νt )) isbounded as ∥∇2ν,θ Vs (νt , θK (νt ))∥ ≤ Hℓ2 Lr B, where6Submission and Formatting Instructions for ICML 2023

3.5

Utility Convergence

Goal Position1 Convergence

Inner Reward (Max) Convergence

0

4.0

3.0
2.5
2.0

Alignment1
Alignment2
Alignment3
Alignment4

1.5
1.0
0

10

1

2

3

Alignment1
Alignment2
Alignment3
Alignment4

4
20

30

Outer Iterations

40

0

50

10

20

30

Outer Iterations

(a)

40

50

Alignment1
Alignment2
Alignment3
Alignment4

40
60
80
100
120
140
0

10

20

30

Outer Iterations

(b)

40

50

(c)

Figure 1. (a) describes the convergence of the upper-level optimization variable with the designer’s alignment levels. Alignments 1, 2, 3,
and 4 used λ = 1, 10, 15, 30, respectively. It is evident that our Algorithm aligns efficiently with the total outer level utility i.e consideringboth the aspects of outer and inner utilities for the designer. Figure 1(b) shows that when weightage on the cost-utility term in the outerobjective is less, the cost naturally increases with outer iterations (shown as a negative reward). Figure 1(c) shows that for alignment
configurations with low weights on the cost-utility part of the outer objective, inner rewards are achieved is more and vice-versa otherwise.
Overall it demonstrates that our algorithm efficiently aligns to the broader utilities

The proof of Lemma 4.3 is in Appendix 12.6.

5. Experimental Evaluations

Lemma 4.4. Under Assumptions 1-(4), the lower level iterates of Algorithm 1 satisfies

To demonstrate our algorithm’s performance, we use a simple grid-world environment where the agent’s objective is toreach a goal at position (Gx , Gy ). The agent starts at (0, 0)
and receives a reward rν based on its speed in reaching thegoal. The goal’s positionp
affects an outer utility/cost Rutil ,which incurs a cost of − (Tx − Gx )2 + (Ty − Gy )2 as it
moves away from the target position (Tx , Ty ). We formulate
this as a bilevel optimization problem with an overall cost
Router = Rinner + λRutil , where λ is a scalar weighting
term. Smaller λ values prioritize inner utility, causing thegoal to move closer to the agent’s start position. Increasing
λ gives more weight to outer utility, bringing the goal closer
to the target position. Fig.1a shows goal position changes
under different alignments, while Fig.1b and Fig.1c depict
the convergence of outer and inner utility to the alignments.
Alignments 1, 2, 3, and 4 correspond to λ = 1, 10, 15, 30
respectively. Lower λ values lead to a goal closer to theagent and higher inner utility, while higher λ values keep thegoal near the target with higher outer utility. The ablation
study validates the algorithm’s alignment performance.

∥θK (νt ) − θ∗ (ν)∥2 ≤η K L6Z,µ(23)where, Z := maxν ∥θ0 − θ∗ (ν)∥2 , η := 1 − α3 , α3 =22
αℓ (1 − αℓ2L6 ) µ2 , L6 = L5 HL
2 + L5 , L5 = Hl RL1 , µ
is the PL constant, and K denotes the number of lowerlevel iterations, and policy gradient step-size satisfies αℓ <2/ min(Hℓ L2 , µ), with µ as the PL constant (Assumption
4).
The proof of Lemma 4.4 is provided in Appendix 12.7. The
proof relies on the assumption that the Value function satisfies a Polyak Łojasiewicz (PL) condition under appropriate
policy parametrization.
Theorem 1. Under Assumptions 1-(4), for the proposed
Algorithm 1, it holds thatT

1X
G0 − G∗ηδ2 L6 Z∥∇ν G(νt , θ∗ (νt ))∥2 ≤
+
T t=1
δ1 T
T δ1 µ(1 − η)(24)

6. Conclusions and Future Work
Potentially misaligned agents pose a severe risk to society;thus making AI alignment at the forefront of research of
the current times (Liu et al., 2022). However, an efficientsolution to such an alignment problem depends upon two
major factors 1) precise Evaluation of the agent’s policy and2) subsequent reward design, and failure to adhere to anyof these can cause potentially misaligned agents or agentswith compromised performance. To deal with the above
entanglement, we formulate the alignment problem in a
bilevel optimization framework and characterized a precise

where G0 := G(ν0 , θ∗ (ν0 )) and G∗ denotes
the global op
1 − 2c11 − Lg αutimum of the outer objective, δ1 = αuand δ2 = αu c21 + Lg αu , c1 is a positive constant definedin eqn. (33), and the step-size range of satisfies αu < 1/Lg ,with Lg as in Assumption 1 and αℓ as stated in Lemma 4.4.
In Theorem 1, we note that we achieve a final rate ofO(1/T ), which matches with bilevel optimization for nonconvex outer objective (Ghadimi & Wang, 2018a) when thelower-level satisfies PL condition.
7

Submission and Formatting Instructions for ICML 2023

evaluation by approximately estimating the gradient of the
optimal policy w.r.t to the reward parametrization, whichis crucial for alignment. Finally, through our design andanalysis of the algorithm, we show convergence guarantees
to the broader goals with a rate O(1/T ), which is a first step
towards provable policy alignment research.

Bertsimas, D., Iancu, D. A., and Parrilo, P. A. Optimality ofaffine policies in multistage robust optimization. Mathematics of Operations Research, 35(2):363–394, 2010.
Bhandari, J. and Russo, D. On the linear convergence ofpolicy gradient methods for finite mdps. In International
Conference on Artificial Intelligence and Statistics, pp.
2386–2394. PMLR, 2021.

References

Boyd, S. P. and Vandenberghe, L. Convex optimization.
Cambridge university press, 2004.

Agarwal, A., Kakade, S. M., Lee, J. D., and Mahajan, G. Optimality and approximation with policy gradient methods
in markov decision processes. In Conference on Learning
Theory, pp. 64–66. PMLR, 2020.

Bracken, J. and McGill, J. T. Mathematical programs with
optimization problems in the constraints. Operations
research, 21(1):37–44, 1973.

Akhtar, Z., Bedi, A. S., Thomdapu, S. T., and Rajawat, K.
Projection-free stochastic bi-level optimization. IEEE
Transactions on Signal Processing, 70:6332–6347, 2022.

Brown, D. S., Goo, W., Nagarajan, P., and Niekum, S. Extrapolating beyond suboptimal demonstrations via inverse
reinforcement learning from observations, 2019.

Ardon, L., Vadori, N., Spooner, T., Xu, M., Vann, J., andGanesh, S. Towards a fully rl-based market simulator. In
Proceedings of the Second ACM International Conference
on AI in Finance, pp. 1–9, 2021.

Buehler, H., Gonon, L., Teichmann, J., and Wood, B. Deep
hedging. Quantitative Finance, 19(8):1271–1291, 2019.
Caplin, A. and Dean, M. Revealed preference, rational
inattention, and costly information acquisition. American
Economic Review, 105(7):2183–2203, 2015.

Arora, S. and Doshi, P. A survey of inverse reinforcement
learning: Challenges, methods and progress, 2020.
Arora, S. and Doshi, P. A survey of inverse reinforcement
learning: Challenges, methods and progress. Artificial
Intelligence, 297:103500, 2021.

Chen, P.-Y., Tielman, M. L., Heylen, D. K. J., Jonker, C. M.,
and van Riemsdijk, M. B. Ai alignment dialogues: Aninteractive approach to ai alignment in support agents,
2023a.

Bai, Y., Jin, C., Wang, H., and Xiong, C. Sample-efficientlearning of stackelberg equilibria in general-sum games.
Advances in Neural Information Processing Systems, 34:
25799–25811, 2021.

Chen, S., Wu, J., Wu, Y., and Yang, Z. Learning to incentivize information acquisition: Proper scoring rules meetprincipal-agent model. arXiv preprint arXiv:2303.08613,
2023b.
Chen, T., Sun, Y., and Yin, W. Closing the gap: Tighter analysis of alternating stochastic gradient methods for bilevelproblems. Advances in Neural Information Processing
Systems, 34:25294–25307, 2021.

Bai, Y., Jones, A., Ndousse, K., Askell, A., Chen, A., DasSarma, N., Drain, D., Fort, S., Ganguli, D., Henighan,T., Joseph, N., Kadavath, S., Kernion, J., Conerly, T.,
El-Showk, S., Elhage, N., Hatfield-Dodds, Z., Hernandez, D., Hume, T., Johnston, S., Kravec, S., Lovitt, L.,
Nanda, N., Olsson, C., Amodei, D., Brown, T., Clark, J.,
McCandlish, S., Olah, C., Mann, B., and Kaplan, J. Training a helpful and harmless assistant with reinforcement
learning from human feedback, 2022.

Chen, Y., Zhang, X., Zhang, K., Wang, M., and Zhu, X.
Byzantine-robust online and offline distributed reinforcement learning, 2022.
Eisenhardt, K. M. Agency theory: An assessment andreview. Academy of management review, 14(1):57–74,
1989.

Balcan, M.-F., Daniely, A., Mehta, R., Urner, R., and Vazirani, V. V. Learning economic parameters from revealed
preferences. In Web and Internet Economics: 10th International Conference, WINE 2014, Beijing, China, December 14-17, 2014. Proceedings 10, pp. 338–353. Springer,
2014.

Fiez, T., Chasnov, B., and Ratliff, L. J. Convergence oflearning dynamics in stackelberg games. arXiv preprint
arXiv:1906.01217, 2019.
Fiez, T., Chasnov, B., and Ratliff, L. Implicit learning dynamics in stackelberg games: Equilibria characterization,
convergence analysis, and empirical study. In International Conference on Machine Learning, pp. 3133–3144.
PMLR, 2020.

Bertsimas, D. and Caramanis, C. Finite adaptability inmultistage linear optimization. IEEE Transactions on
Automatic Control, 55(12):2751–2766, 2010.
8

Submission and Formatting Instructions for ICML 2023

Fürnkranz, J., Hüllermeier, E., Cheng, W., and Park, S.H. Preference-based reinforcement learning: a formal
framework and a policy iteration algorithm. Machinelearning, 89:123–156, 2012.

Jabbari, S., Joseph, M., Kearns, M., Morgenstern, J., andRoth, A. Fairness in reinforcement learning. In International conference on machine learning, pp. 1617–1626.
PMLR, 2017.

Ghadimi, S. and Wang, M. Approximation methods forbilevel programming. arXiv preprint arXiv:1802.02246,
2018a.

Ji, K. and Liang, Y. Lower bounds and accelerated algorithms for bilevel optimization. Journal of Machine
Learning Research, 23:1–56, 2022.

Ghadimi, S. and Wang, M. Approximation methods forbilevel programming, 2018b. URL https://arxiv.
org/abs/1802.02246.

Kang, B., Jie, Z., and Feng, J. Policy optimization withdemonstrations. In Dy, J. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine
Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2469–2478. PMLR, 10–15 Jul 2018.
URL https://proceedings.mlr.press/v80/
kang18a.html.

Ghasemipour, S. K. S., Zemel, R., and Gu, S. A divergence
minimization perspective on imitation learning methods,
2019.
Goktas, D., Zhao, S., and Greenwald, A. Zero-sum stochastic stackelberg games. Advances in Neural Information
Processing Systems, 35:11658–11672, 2022.

Khanduri, P., Zeng, S., Hong, M., Wai, H.-T., Wang, Z., andYang, Z. A near-optimal algorithm for stochastic bilevel
optimization via double-momentum. Advances in neural
information processing systems, 34:30271–30283, 2021.

Green, J. R. and Stokey, N. L. A comparison of tournaments
and contracts. Journal of Political Economy, 91(3):349–
364, 1983.

Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline reinforcement learning: Tutorial, review, and perspectives onopen problems, 2020.

Hansen, P., Jaumard, B., and Savard, G. New branch-andbound rules for linear bilevel programming. SIAM Journal on scientific and Statistical Computing, 13(5):1194–
1217, 1992.

Li, J., Gu, B., and Huang, H. A fully single loop algorithm for bilevel optimization without hessian inverse. In
Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 7426–7434, 2022.

Hill, E. W., Bardoscia, M., and Turrell, A. E. Solving
heterogeneous general equilibrium economic models withdeep reinforcement learning. ArXiv, abs/2103.16977,
2021.

Lillicrap, T. P., Hunt, J. J., Pritzel, A., Heess, N., Erez,
T., Tassa, Y., Silver, D., and Wierstra, D. Continuous
control with deep reinforcement learning. arXiv preprint
arXiv:1509.02971, 2015.

Ho, J. and Ermon, S. Generative adversarial imitation learning, 2016.

Liu, R., Zhang, G., Feng, X., and Vosoughi, S. Aligning generative language models with human values. In Findings of the Association for Computational Linguistics: NAACL 2022, pp. 241–252, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.findings-naacl.
18. URL https://aclanthology.org/2022.
findings-naacl.18.

Hong, M., Wai, H.-T., Wang, Z., and Yang, Z. A twotimescale framework for bilevel optimization: Complexity analysis and application to actor-critic. arXiv preprint
arXiv:2007.05170, 2020.
Hu, Z., Liang, Y., Zhang, J., Li, Z., and Liu, Y. Inference
aided reinforcement learning for incentive mechanism
design in crowdsourcing. Advances in Neural Information
Processing Systems, 31, 2018.

Lyu, B., Meng, Q., Qiu, S., Wang, Z., Yang, Z., and Jordan, M. I. Learning dynamic mechanisms in unknown
environments: A reinforcement learning approach. arXiv
preprint arXiv:2202.12797, 2022a.

Huang, F., Li, J., Gao, S., and Huang, H. Enhanced bilevel
optimization via bregman distance. Advances in Neural
Information Processing Systems, 35:28928–28939, 2022.

Lyu, B., Wang, Z., Kolar, M., and Yang, Z. Pessimism meets
vcg: Learning dynamic mechanism design via offline
reinforcement learning. In International Conference on
Machine Learning, pp. 14601–14638. PMLR, 2022b.

Hurwicz, L. Mechanism design without games. Advancesin Economic Design, pp. 429–437, 2003.
Isbell, C., Shelton, C. R., Kearns, M., Singh, S., and Stone, P.
A social reinforcement learning agent. In Proceedings ofthe fifth international conference on Autonomous agents,
pp. 377–384, 2001.

Maei, H. R., Szepesvári, C., Bhatnagar, S., Precup, D., Silver, D., and Sutton, R. S. Convergent temporal-difference
learning with arbitrary smooth function approximation.
9Submission and Formatting Instructions for ICML 2023

In Proceedings of the 22nd International Conference on
Neural Information Processing Systems, NIPS’09, pp.
1204–1212, Red Hook, NY, USA, 2009. Curran Associates Inc. ISBN 9781615679119.

Rahwan, I. Society-in-the-loop: programming the algorithmic social contract. Ethics and information technology,
20(1):5–14, 2018.
Roth, A., Ullman, J., and Wu, Z. S. Watch and learn:
Optimizing from revealed preferences feedback. In
Proceedings of the Forty-Eighth Annual ACM Symposium on Theory of Computing, STOC ’16, pp. 949–962,
New York, NY, USA, 2016. Association for Computing Machinery. ISBN 9781450341325. doi: 10.1145/
2897518.2897579. URL https://doi.org/10.
1145/2897518.2897579.

Maheshwari, C., Sasty, S. S., Ratliff, L., and Mazumdar, E.
Convergent first-order methods for bi-level optimizationand stackelberg games. arXiv preprint arXiv:2302.01421,
2023.
Maskin, E. S. Mechanism design: How to implement social
goals. American Economic Review, 98(3):567–576, 2008.
Mei, J., Xiao, C., Szepesvari, C., and Schuurmans, D. On
the global convergence rates of softmax policy gradientmethods. In International Conference on Machine Learning, pp. 6820–6829. PMLR, 2020a.

Saha, A., Pacchiano, A., and Lee, J. Dueling rl: Reinforcement learning with trajectory preferences. In International Conference on Artificial Intelligence and Statistics,
pp. 6263–6289. PMLR, 2023.

Mei, J., Xiao, C., Szepesvari, C., and Schuurmans, D. On
the global convergence rates of softmax policy gradient methods. In III, H. D. and Singh, A. (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 6820–6829. PMLR, 13–
18 Jul 2020b. URL https://proceedings.mlr.
press/v119/mei20b.html.
Mnih, V., Kavukcuoglu, K., Silver, D., Graves, A.,
Antonoglou, I., Wierstra, D., and Riedmiller, M. Playing
atari with deep reinforcement learning. arXiv preprint
arXiv:1312.5602, 2013.

Shi, C., Lu, J., and Zhang, G. An extended kuhn–tucker
approach for linear bilevel programming. Applied Mathematics and Computation, 162(1):51–63, 2005.
Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou,
I., Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M.,
Bolton, A., et al. Mastering the game of Go without
human knowledge. Nature, 550(7676):354–359, 2017.
Steck, H., Baltrunas, L., Elahi, E., Liang, D., Raimond, Y.,
and Basilico, J. Deep learning for recommender systems:
A netflix case study. AI Magazine, 42(3):7–18, 2021.

Myerson, R. B. Mechanism design. Springer, 1989.

Stiglitz, J. E. Principal and agent. Springer, 1989.

Ng, A. Y., Harada, D., and Russell, S. Policy invarianceunder reward transformations: Theory and application
to reward shaping. In Icml, volume 99, pp. 278–287.
Citeseer, 1999.

Sun, Z., Shen, Y., Zhou, Q., Zhang, H., Chen, Z., Cox, D.,
Yang, Y., and Gan, C. Principle-driven self-alignmentof language models from scratch with minimal humansupervision, 2023.

Ng, A. Y., Russell, S., et al. Algorithms for inverse reinforcement learning. In Icml, volume 1, pp. 2, 2000.

Sutton, R. S. and Barto, A. G. Reinforcement learning: An
introduction. Cambridge: MIT press, 1998.

Ni, Z. and Paul, S. A multistage game in smart grid security:
A reinforcement learning solution. IEEE transactions onneural networks and learning systems, 30(9):2684–2695,
2019.

Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y.
Policy gradient methods for reinforcement learning withfunction approximation. Advances in neural information
processing systems, 12, 1999.

Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C.,
Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A.,et al. Training language models to follow instructions
with human feedback. Advances in Neural Information
Processing Systems, 35:27730–27744, 2022.

Sutton, R. S., Maei, H. R., Precup, D., Bhatnagar, S.,
Silver, D., Szepesvári, C., and Wiewiora, E. Fast
gradient-descent methods for temporal-difference learning with linear function approximation. In Proceedings of the 26th Annual International Conference on
Machine Learning, ICML ’09, pp. 993–1000, New
York, NY, USA, 2009. Association for Computing
Machinery. ISBN 9781605585161. doi: 10.1145/
1553374.1553501. URL https://doi.org/10.
1145/1553374.1553501.

Pflug, G. C. and Pichler, A. Multistage stochastic optimization, volume 1104. Springer, 2014.
Puterman, M. L. Markov Decision Processes: Discrete
stochastic dynamic programming. John Wiley & Sons,
2014.
10

Submission and Formatting Instructions for ICML 2023

Tang, P. Reinforcement mechanism design. In IJCAI, pp.
5146–5150, 2017.

Yin, M., Wang, M., and Wang, Y.-X. Offline reinforcement
learning with differentiable function approximation isprovably efficient, 2022.

Teo, Y. M., Luong, B. L., and Szabo, C. Formalization of
emergence in multi-agent systems. In Proceedings of the
1st ACM SIGSIM Conference on Principles of Advanced
Discrete Simulation, pp. 231–240, 2013.

Zhang, K., Koppel, A., Zhu, H., and Basar, T. Global convergence of policy gradient methods to (almost) locally
optimal policies. SIAM Journal on Control and Optimization, 58(6):3586–3612, 2020.

Torabi, F., Warnell, G., and Stone, P. Behavioral cloning
from observation. In Proceedings of the 27th International Joint Conference on Artificial Intelligence, pp.
4950–4957, 2018.

Zhao, G., Zhu, B., Jiao, J., and Jordan, M. I. Online learningin stackelberg games with an omniscient follower. arXiv
preprint arXiv:2301.11518, 2023.

Von Stackelberg, H. Market structure and equilibrium.
Springer Science & Business Media, 2010.

Zhong, H., Yang, Z., Wang, Z., and Jordan, M. I. Can
reinforcement learning find stackelberg-nash equilibria
in general-sum markov games with myopic followers?
arXiv preprint arXiv:2112.13521, 2021.

Vu, Q.-L., Alumbaugh, Z., Ching, R., Ding, Q., Mahajan,
A., Chasnov, B., Burden, S., and Ratliff, L. J. Stackelberg
policy gradient: Evaluating the performance of leadersand followers. In ICLR 2022 Workshop on Gamificationand Multiagent Solutions.

Zhu, B., Bates, S., Yang, Z., Wang, Y., Jiao, J., and Jordan,
M. I. The sample complexity of online contract design.
arXiv preprint arXiv:2211.05732, 2022.
Zhu, B., Jiao, J., and Jordan, M. I. Principled reinforcement learning with human feedback from pairwise or
k-wise comparisons, 2023. URL https://arxiv.
org/abs/2301.11270.

Wen, C., Lin, J., Darrell, T., Jayaraman, D., and Gao, Y.
Fighting copycat agents in behavioral cloning from observation histories. Advances in Neural Information Processing Systems, 33:2564–2575, 2020.

Ziebart, B. D., Maas, A., Bagnell, J. A., and Dey, A. K.
Maximum entropy inverse reinforcement learning. In
Proceedings of the 23rd National Conference on Artificial
Intelligence - Volume 3, AAAI’08, pp. 1433–1438. AAAI
Press, 2008. ISBN 9781577353683.

Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Reinforcement learning, pp. 5–32, 1992.
Wirth, C., Akrour, R., Neumann, G., Fürnkranz, J., et al. A
survey of preference-based reinforcement learning methods. Journal of Machine Learning Research, 18(136):
1–46, 2017.
Wolf, Y., Wies, N., Levine, Y., and Shashua, A. Fundamental limitations of alignment in large language models,
2023.
Xiao, H., Herman, M., Wagner, J., Ziesche, S., Etesami, J.,
and Linh, T. H. Wasserstein adversarial imitation learning,
2019.
Xu, Y., Wang, R., Yang, L., Singh, A., and Dubrawski, A.
Preference-based reinforcement learning with finite-timeguarantees. Advances in Neural Information Processing
Systems, 33:18784–18794, 2020.
Yang, J., Ji, K., and Liang, Y. Provably faster algorithms
for bilevel optimization. Advances in Neural Information
Processing Systems, 34:13670–13682, 2021.
Yang, S., Nachum, O., Du, Y., Wei, J., Abbeel, P., andSchuurmans, D. Foundation models for decision making:Problems, methods, and opportunities, 2023.
11

Submission and Formatting Instructions for ICML 2023

Appendix
7. Notations
We collectively describe the notations used in this work in Table 1.
Table 1:
Notations

Details

S, A

State space, action space(s, a)

State-action pair′P(s |s, a)Transition kernelrν (s, a)Reward function parameterized by designer’s parameters ν ∈ Rn

πθ(ν) (·|s)

Policy parameterized by θ(ν) ∈ Rd for design parameters ν

Vs (ν, θ(ν))

Inner objective - Value function for state s at outer parameter ν and policy
parameter θ(ν)

G(ν, θ∗ (ν))

Outer objective

8. Related Work
Bilevel Optimization. Multi-stage optimization has a long-history in optimization and operations research, both for
deterministic (Bertsimas & Caramanis, 2010) and stochastic objectives (Pflug & Pichler, 2014). In general, these problems
exhibit NP-hardness if the objectives are general non-convex functions. Therefore, much work in this area restricts focus toclasses of problems, such as affine (Bertsimas & Caramanis, 2010; Bertsimas et al., 2010). From both the mathematicaloptimization community (Bracken & McGill, 1973) and algorithmic game theory (Von Stackelberg, 2010), extensive interest
has been given to specifically two-stage problems. A non-exhaustive list contains the following works (Ghadimi & Wang,
2018b; Yang et al., 2021; Khanduri et al., 2021; Li et al., 2022; Ji & Liang, 2022; Huang et al., 2022). Predominately, one
these works do not consider that either stage is a MDP, and therefore while some of the algorithmic strategies are potentiallygeneralizable to the RL agent alignment problem, they are not directly comparable to the problem studied here.
Algorithms for Stackleberg Games. Algorithmic methods for Sackleberg games are a distinct line of inquiry that develops
methods to reach the Stackleberg-Nash equilibrium of a game, which have received significant attention in recent years.
Beginning with the static Stackleberg game setting, conditions for gradient play to achieve local equilibria have been
established in (Fiez et al., 2019). Follow on work studied conditions for gradient play to converge without convexity, butdoes not allow for MDP/trajectory dependence of objective functions at either stage (Maheshwari et al., 2023). On the other
hand, the access to information structures have gradually been relaxed to allow bandit feedback (Bai et al., 2021) and linear
MDPs (Zhong et al., 2021). Value iteration schemes have been proposed to achieve Stackleberg-Nash equilibria as well
(Goktas et al., 2022), although it is unclear how to generalize them to handle general policy parameterization in a scalable
manner. Most similar to our work are those that develop implicit function-theorem based gradient play (Fiez et al., 2020; Vu
et al.); however, there are no performance certificates for these approaches.
Mechanism Design. In this line of research, one studies the interrelationship between the incentives of an individual
economic actor and their macro-level behavior at the level of a social welfare objective. This literature can be traced back to(Myerson, 1989; Hurwicz, 2003; Maskin, 2008), and typically poses the problem as one that does not involve sequential
interactions. More recently, efforts to cast the evolution of the outer-stage which quantifies social welfare or ethnicalconsiderations as a sequential process, i.e., an MDP, have been considered (Tang, 2017; Hu et al., 2018). In these works,
agents’ behavior is treated as fixed and determining of the state transition dynamics, which gives rise to a distinct subclassof policy optimization problems (Lyu et al., 2022a;b).
Principal-Agent Problem. Contract theory also formalizes the notion of an agent’s interaction with a rule setter (Green &Stokey, 1983; Eisenhardt, 1989; Stiglitz, 1989). The agent is the individual who may choose to enroll in a contract, and theprincipal is the creator of the rules or incentive structure that determines the incentive structure for the agent’s decision.
Formalizing the question of how to design a contract and agent behavior through MDP machinery and Stackleberg games
12Submission and Formatting Instructions for ICML 2023

Algorithm 2 Policy Alignment with Reward Design (conceptual)
1: Input: Reward parametrization ν0 , optimal policy parameter θ ∗ (ν0 ) (oracle), upper-level step size α > 0
2: for all t = 0, 1, 2, ..., T − 1 do
3:Update the parameter νt as

νt+1 = νt − α∇ν G[νt , θ∗ (νt )](25)

where ∇ν G[ν, θ∗ (νt )] is defined in (13).
4: end for
5: Output: νT

has been the subject of recent efforts, especially (Zhu et al., 2022; Chen et al., 2023b).
RL with Preferences. Revealed preferences is a concept in economics that states that humans in many cases do notknow what they prefer until they are exposed to examples (Caplin & Dean, 2015). This idea gained traction as a itprovided a substantive basis for querying humans and elucidating feedback on decision-making problems with metrics
whose quantification is not obvious, such as trustworthiness or fairness. Efforts to incorporate pairwise comparison into
reinforcement learning, especially as a mechanism to incorporate preference information, has been heavily studied in recent
years – see (Wirth et al., 2017) for a survey. A non-exhaustive list of works along these lines is (Isbell et al., 2001; Fürnkranz
et al., 2012; Roth et al., 2016; Hill et al., 2021; Wirth et al., 2017; Zhu et al., 2023; Xu et al., 2020; Saha et al., 2023).
Inverse RL and Behavioral Cloning. Implicitly contained in the RL with preferences framework is the assumption thathumans should design a reward function to drive an agent towards the correct behavior through its policy optimization process.
This question of reward design has also been studied through an alternative lens in which one provides demonstrations
or trajectories, and seeks to fit a reward function to represent this information succinctly. This class of problems, broadly
referred to as inverse RL (Ng et al., 2000), has a long history in experimentally driven RL research – see (Arora & Doshi,2021) for a survey. Rather than design a reward function as an intermediary between demonstrations and policy optimization,
behavioral cloning directly seeks to mimic the behavior of demonstration information (Torabi et al., 2018; Wen et al., 2020).
In doing so, it is able to sidestep some of the questions of whether a reward function can be well-posed for a given collectionof trajectories.

9. Upper & Lower Level Gradient Computations
Upper-level Gradient Estimation To estimate the gradient of the upper-level objective in equation (8), we require
information about θ∗ (ν) (issue (i)) which is not available in general unless the inner-level objective has a closed-form
solution. Hence, we approximate θ∗ (νt ) with θK (νt ) i.e., running K-step policy gradient steps at the lower-level as
"
#
HX
u −1
K
K
K
e
f
∇ν G(νt , θ (νt )) = EP (τ ;θK (νt )) U (τ ) ·[M (νt )∇θ fh (θ (νt ))] + ∇ν Z(ν).

(26)

t=0

fK (νt ) = −∇2 Vs (νt , θK (νt ))∇2 Vs (νt , θK (νt ))−1 . This approximate evaluation of (13) with a current policy
where Mv,θθestimate θK (νt ) in lieu of θ∗ (νt ) constitutes our effort to address issue (iii).
Lower-level objective gradient, Jacobian, and Hessian estimation: Next, we derive the exact gradients for the inner-levelobjective and, subsequently the Hessian and mixed hessian terms for our algorithmic description (issue (ii)). Let us firstwrite down the gradient of lower-level objectives using the Policy Gradient Theorem (Williams, 1992; Sutton et al., 1999) as∇θ Vs (νt , θK (νt )) = EP (τ ;θK (νt ))" H −1
ℓ
X



H
ℓ −1
X

γ h rνt (sh , ah ) 

j=0

h=0

13

#
∇θ log πθK (νt ) (aj |sj ) ,(27)

Submission and Formatting Instructions for ICML 2023

where Hℓ is the horizon or the episode length. Similarly, the Hessian of the inner objective is:#

γ h rνt (sh , ah ) ∇2θ log πθK (νt ) (aj |sj ) ,

(28)

"H −1
 HX
#
ℓ
ℓ −1
X
h
T
γ ∇ν rνt (sh , ah )[∇θ log πθK (νt ) (aj |sj )].

(29)∇2θ Vs (νt , θK (νt )) = EP (τ ;θK (νt ))" H −1
ℓ
X

H
ℓ −1
X
j=0

h=0

Finally, we can write the mixed second-order Jacobian matrix as∇2ν,θ Vs (νt , θK (νt )) = EP (τ ;θK (νt ))

j=0

h=0

Now, utilizing the expressions in (26)-(29), we summarize the proposed steps in Algorithm 1. We note that we denotethe number of inner-level iterates by K for better exposition. We will choose it as a function of outer iterations t in theconvergence analysis next. Without loss of generality, we have written the algorithm updates as if we are minimizing theupper and lower level objectives. Before shifting to analyzing the convergence behavior of Algorithm 1, we close with aremark.

10. Proof of Theorem 1
Proof. We begin by the smoothness assumption in the upper-level objective (cf. Assumption 14), which implies thatG[νt+1 , θ∗ (νt+1 )] ≤ G(νt , θ∗ (νt )) + ⟨∇ν G(νt , θ∗ (νt )), νt+1 − νt ⟩ +Lg∥νt+1 − νt ∥2 .
2

(30)

From the update of outer parameter νt+1 (cf. (26)), we holds thate ν G(νt , θK (νt ))⟩G(νt+1 , θ∗ (νt+1 )) ≤G(νt , θ∗ (νt )) + ⟨∇ν G(νt , θ∗ (νt )), −αu ∇+Lg αu22(31)e ν G(νt , θK (νt ))∥2 .
∥∇We add subtract the original gradient ∇ν G(νt , θ∗ (νt )) [cf. (13)] in (31) as follows
G(νt+1 , θ∗ (νt+1 )) ≤G(νt , θ∗ (νt )) + ⟨∇ν G(νt , θ∗ (νt )), −αu ∇ν G(νt , θ∗ (νt ))⟩e ν G(νt , θK (νt ))⟩+ αu ⟨∇ν G(νt , θ∗ (νt )), ∇ν G(νt , θ∗ (νt )) − ∇Lg αu2e ν G(νt , θK (νt )) − ∇ν G(νt , θ∗ (νt ))∥2∥∇ν G(νt , θ∗ (νt )) + ∇2=G(νt , θ∗ (νt )) − αu ∥∇ν G(νt , θ∗ (νt ))∥2e ν G(νt , θK (νt ))⟩+ αu ⟨∇ν G(νt , θ∗ (νt )), ∇ν G(νt , θ∗ (νt )) − ∇

(32)

+

+

(33)

Lg αu2e ν G(νt , θK (νt )) − ∇ν G(νt , θ∗ (νt ))∥2 .
∥∇ν G(νt , θ∗ (νt )) + ∇2

Using Peter-Paul inequality for the third term on the right hand side of (33), we get
G(νt+1 , θ∗ (νt+1 )) ≤G(νt , θ∗ (νt )) − αu ∥∇ν G(νt , θ∗ (νt ))∥2 ++αu∥∇ν G(νt , θ∗ (νt ))∥2
2c1

αu c1
e ν G(νt , θK (νt ))∥2∥∇ν G(νt , θ∗ (νt )) − ∇2
Lg αu2e ν G(νt , θK (νt )) − ∇ν G(νt , θ∗ (νt ))∥2 .
+∥∇ν G(νt , θ∗ (νt )) + ∇
2
14

(34)

Submission and Formatting Instructions for ICML 2023

where c1 ≥ 0. Next, after grouping the terms, we get1G(νt+1 , θ∗ (νt+1 )) ≤G(νt , θ∗ (νt )) − αu 1 −∥∇ν G(νt , θ∗ (νt ))∥2
2c1
αu c1
e ν G(νt , θK (νt ))∥2+∥∇ν G(νt , θ∗ (νt )) − ∇2e ν G(νt , θK (νt )) − ∇ν G(νt , θ∗ (νt ))∥2
+ Lg αu2 ∥∇ν G(νt , θ∗ (νt ))∥2 + Lg αu2 ∥∇1=G(νt , θ∗ (νt )) − αu 1 −− Lg αu ∥∇ν G(νt , θ∗ (νt ))∥2
2c1

c
1
e ν G(νt , θK (νt ))∥2 ,+ Lg αu ∥∇ν G(νt , θ∗ (νt )) − ∇
+ αu
2

(35)

where we use the inequality ∥a + b∥2 ≤ 2∥a∥2 + 2∥b∥2 , followed by algebraic operations to get the final expression ofe ν G(νt , θK (νt ))∥2 from equation (35). Let us start byequation (35). Next, we analyze the term ∥∇ν G(νt , θ∗ (νt )) − ∇∗econsidering the explicit expressions of ∇ν G(ν, θ (ν)) and ∇ν G(νt , θK (νt )) in equations (13) and (26) ase ν G(ν, θK (ν)) = Eτ ∼P (τ ;θ∗ (ν)) [ϕ1 (τ )] − Eτ ∼P (τ ;θK (ν)) [ϕ2 (τ )],∇ν G(ν, θ∗ (ν)) − ∇(36)

where we define
ϕ1 (τ ) =U (τ ) ·Hu
X[−∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θ∗ (ν))−1 ∇θ fh (θ∗ (ν))],(37)[−∇2v,θ Vs (ν, θK (ν))∇2θ Vs (ν, θK (ν))−1 ∇θ fh (θK (ν))].

(38)

t=0

and ϕ2 (τ ) =U (τ ) ·

Hu
X
t=0

Now, we expand the terms in (36) by adding subtracting the term Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ )] in the right hand side as follows
e ν G(ν, θK (ν)) =Eτ ∼P (τ ;θ∗ (ν)) [ϕ1 (τ )] − Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ )]∇ν G(ν, θ∗ (ν)) − ∇(39)+ Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ )] − Eτ ∼P (τ ;θK (ν)) [ϕ2 (τ )]=Eτ ∼P (τ ;θ∗ (ν)) [ϕ1 (τ )] − Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ )]+ Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ ) − ϕ2 (τ )].

(40)

We note that the first term on the right-hand side of (40) Eτ ∼P (τ ;θ∗ (ν)) [ϕ1 (τ )] − Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ )] ≤supϕ Eτ ∼P (τ ;θ∗ (ν)) [ϕ1 (τ )] − Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ )]. This boils down to the standard definition of Integral Probability
Metric(IPM) which, under suitable assumption on the function ϕ can be upper-bounded by Wasserstein and Total Variation
distance. Specifically, we show that function ϕ is Lipschitz with some constant. The second term on the right-hand side of(40) is the expected difference between the two functions. Hence, we can write (40) ase ν G(ν, θK (ν)) =Df (P (τ ; θ∗ (ν)), P (τ ; θK (ν)))∇ν G(ν, θ∗ (ν)) − ∇+ Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ ) − ϕ2 (τ )],(41)

where Df denotes the f-divergence between two distributions. Taking the norm on both sides and from the statements ofLemma 4.1 and Lemma 4.3, we can write 2 2

Hu L 2
∗
K
2
4 2 2
e
∥∇ν G(ν, θ (ν)) − ∇ν G(ν, θ (ν))∥ ≤+ 2Hu ũ γ1 ∥θ∗ (ν) − θK (ν)∥2 .
(42)
2
Utilizing this bound in (35), we get
G(νt+1 ,θ∗ (νt+1 )) − G(νt , θ∗ (νt ))1≤ −αu 1 −− Lg αu ∥∇ν G(νt , θ∗ (νt ))∥22c1
c
  H 2 L2
1
u 2
4 2 2
+ αu
+ Lg αu+ 2Hu ũ γ1 ∥θ∗ (ν) − θK (ν)∥2 ,
2
2
15

(43)

Submission and Formatting Instructions for ICML 2023

where we write thefrom equation (35) final expression for the convergence analysis
 and for simplicity of notations, let’s  Hu2 L22
4 2 2
assume δ1 = αu 1 − 2c11 − Lg αu and δ2 = αu c21 + Lg αu
+
2H
ũ
γ
u
1 , which leads to the simplified version
2of the equation (43)G(νt+1 , θ∗ (νt+1 )) − G(νt , θ∗ (νt )) ≤ −δ1 ∥∇ν G(νt , θ∗ (νt ))∥2 + δ2 ∥θ∗ (νt ) − θK (νt )∥2 .

(44)

From the statement of Lemma 4.4, we can upper bound the above expression asη K L6Z,µ(45)

η t+1 L6
Z.
µ

(46)

G(νt+1 , θ∗ (νt+1 )) − G(νt , θ∗ (νt )) ≤ −δ1 ∥∇ν G(νt , θ∗ (νt ))∥2 + δ2
where we know η ∈ (0, 1). Next, we select K = t + 1 to obtain
G(νt+1 , θ∗ (νt+1 )) − G(νt , θ∗ (νt )) ≤ −δ1 ∥∇ν G(νt , θ∗ (νt ))∥2 + δ2
Taking the summation over t = 0 to T − 1 on both sides, we get
G(νT , θ∗ (νT ) − G(ν0 , θ∗ (ν0 )) ≤ −δ1

T
−1
X

T −1

∥∇ν G(νt , θ∗ (νt ))∥2 +t=0

δ2 L6 Z X t+1
η
µ
t=0

(47)

After rearranging the terms, we get
T−1X∥∇ν G(νt , θ∗ (νt ))∥2 ≤

t=0

≤

T −1
G(ν0 , θ∗ (ν0 )) − G(νT , θ∗ (νT ) ηδ2 L6 Z X tη+δ1δ1 µ t=0

G(ν0 , θ∗ (ν0 )) − G(νT , θ∗ (νT )ηδ2 L6 Z+.
δ1
δ1 µ(1 − η)(48)

Let us denote G0 := G(ν0 , θ∗ (ν0 )) and upper bound −G(νT , θ∗ (νT ) ≤ −G∗ where G∗ denotes the global optimal valueof the outer objective. After dividing both sides in (48) by T , we get
T −1

1 X
G0 − G∗ηδ2 L6 Z∥∇ν G(νt , θ∗ (νt ))∥2 ≤
+
.
T t=0
δ1 T
T δ1 µ(1 − η)(49)

11. Proof of All Lemmas 4.1-4.4
11.1. Proof of Lemma 4.1
Proof. The probability distribution of the trajectory τ = {sh , ah }H
h=1 is given by
P (τ ; θ∗ (ν)) = ρ(s0 )HYπθ∗ (ν) (ah |sh )P(sh+1 |sh , ah ).

(50)

h=1

Similarly, we can derive an equivalent expression for the probability of trajectory induced by the policy πθK (ν) by replacingθ∗ (ν) with θK (ν). Here, P(sh+1 |sh , ah ) is the transition probability which remains the same for both and ρ(s0 ) is the initialstate distribution. Next, the f-divergence between the two distributions Df (P (τ ; θ∗ (ν)), P (τ ; θK (ν))) can be written as
Df (P (τ ; θ∗ (ν)), P (τ ; θK (ν))) ≤ Df (P (τ ; θ∗ (ν)), P (τ ; β)) + Df (P (τ ; β), P (τ ; θK (ν))),
|
{z
} |
{z
}
I

(51)

II

which holds by triangle inequality (of f-divergences). P (τ ; β) represents the trajectory probability induced by anotherhybrid policy πβ (·|s) which executes the action based on the policy πθK (ν) (·|s) for the first time-step and then follows the16Submission and Formatting Instructions for ICML 2023

policy πθ∗ (ν) (·|s) for subsequent timesteps. Now, we focus on term I in (51), we get
Df (P (τ ;θ∗ (ν)), P (τ ; β))=X

P (τ ; β)f

τ

=

X

=

!
(52)

P (τ ; β))f

P (τ ; β))f!
πθ∗ (ν) (a1 |s0 ),πθK (ν) (a1 |s0 )τρ(s0 )πθ∗ (ν) (a1 |s0 )P(s2 |s1 , a1 )QH

h=2 πθ ∗ (ν) (ah |sh )P(sh+1 |sh , ah )QHρ(s0 )πθK (ν) (a1 |s0 )P(s2 |s1 , a1 ) h=2 πθ∗ (ν) (ah |sh )P(sh+1 |sh , ah )

τX

P (τ ; θ∗ (ν))
P (τ ; β)!

where first we expand upon the definition of the trajectory distribution induced by both policies and get the final expressionof the equation (52). By expanding the term P (τ ; β) in (52), we obtain
Df (P (τ ; θ∗ (ν)), P (τ ; β)) =X

ρ(s0 )

s0

=

X
a1

X

ρ(s)

X

s

πθK (ν) (a|s)fa= Eρ(s)Xθ (ν) (a1 |s1 )XπθK (ν) (a1 |s0 ) π ∗ (a|s) s1πθK (ν) (a1 |s0 )fπθK (ν) (a|s)f

a

π ∗

P(s1 |s0 , a1 ) · · ·θ (ν)πθK (ν) (a|s)
 π ∗ (a|s) θ (ν)πθK (ν) (a|s)= Eρ(s) [Df (πθ∗ (ν) (a|s), πθK (ν) (a|s))],(53)

where, in the first equation we expand upon the sum over all trajectories with the occupancy distribution over states andactions, and replacing with f-divergence, we get the final expression.
Next, we expand similarly for the term II in (51) and expand as
Df (P (τ ; β), P (τ ; θK (ν)))=

X

K

P (τ ; θ (ν)))fτ

P (τ ; β)P (τ ; θK (ν))!

ρ(s0 )πθK (ν) (a1 |s0 )P(s2 |s1 , a1 )QH

h=2 πθ ∗ (ν) (ah |sh )P(sh+1 |sh , ah )=P (τ ; θK (ν)))fQHρ0 (s0 )πθK (ν) (a1 |s0 )P(s2 |s1 , a1 ) h=2 πθK (ν) (ah |sh )P(sh+1 |sh , ah )
τ
!
QH
X
X
X
h=2 πθ ∗ (ν) (ah |sh )P(sh+1 |sh , ah )=ρ(s0 )πθK (ν) (a1 |s0 )P(s1 |s0 , a0 ) · · · f QH
s0a1s1
h=2 πθ K (ν) (ah |sh )P(sh+1 |sh , ah )

X

=

X

ρ(s0 )

X

s0

a1

X

ρ(s0 )X=

s0πθK (ν) (a1 |s0 )XP (τ1 ; θK (ν),τ1!

P (τ1 ; θ∗ (ν)
P (τ1 ; θK (ν))πθK (ν) (a1 |s0 )Df (P (τ1 ; θ∗ (ν), P (τ1 ; θK (ν)),

(54)

a1

where we expand the trajectory distribution induced by the policies and subsequently express as the ratio of the probabilityof trajectories wrt τ1 , we get the final expression. Now, we expand upon the f-divergence of the trajectory τ1 distribution as
Df (P (τ ; β), P (τ ; θK (ν))) =X

ρ(s0 )

s0

≤

X
s0

X

πθK (ν) (a1 |s0 )Df (P (τ1 ; θ∗ (ν), P (τ1 ; θK (ν))

a1

ρ(s0 )XπθK (ν) (a1 |s0 ) Df (P (τ1 ; θ∗ (ν)), P (τ1 ; β))a1+ Df (P (τ1 ; β), P (τ1 ; θK (ν))) ,
17

(55)

Submission and Formatting Instructions for ICML 2023

where using the triangle inequality and get back the similar form with which we had started in equation (51) similar to term
I and term II. Here, similarly continuing this expansion, we finally get∗

K

Df (P (τ ; θ (ν)), P (τ ; θ (ν))) ≤H−1
X

Es∼ρP K
θ

h=0

(ν)

(s) Df (πθ ∗ (ν) (a|s), πθ K (ν) (a|s))≤ HEs∼ρθK (ν) (s) Df (πθ∗ (ν) (a|s), πθK (ν) (a|s))XρθK (ν) (s)Df (πθ∗ (ν) (a|s), πθK (ν) (a|s))≤Hs≤ HDf (πθ∗ (ν) (a′ |s′ ), πθK (ν) (a′ |s′ )),(56)

where, we upper bound the first equation by the total number of timesteps or horizon length H of the trajectory andsubsequently upper-bound the divergence by the state (s, a) pair for which the Df (πθ∗ (ν) (a|s), πθK (ν) (a|s)) is maximum
and is given as (s′ , a′ ). Next, in (56), by considering the total variation as the f-divergence and expanding using definition
with countable measures to obtain
Df (P (τ ; θ∗ (ν)), P (τ ; θK (ν))) ≤ HDT V (πθ∗ (ν) (a|s′ ), πθK (ν) (a|s′ ))H≤ ∥πθ∗ (ν) (a|s) − πθK (ν) (a|s)∥1
2
HL2 ∗∥θ (ν) − θK (ν)∥,
≤
2

(57)

where, we use the Lipschitz assumption (cf. Assumption 3) on the policy parameter to get the final expression of equation(57). We note that the result holds for any general horizon length H.
11.2. Proof of Lemma 4.2
11.2.1. P ROOF OF L EMMA 4.2 S TATEMENT ( I )Proof. We start with the definition from (29)∇2ν,θ Vs (νt , θK (νt )) =X

P (τ ; θK (νt ))

" H
ℓ
X

τ

≤

" H
ℓ
X


T #
Hℓ
X
γ h · ∇ν rνt (sh , ah ) · ∇θ log πθK (νt ) (aj |sj )
j=0

h=0


γ h · ∇ν rνt (s′h , a′h ) · 

=

T #
X

∇θ log πθK (νt ) (a′j |s′j )


γ h · ∇ν rνt (s′h , a′h ) · HℓX

P (τ ; θK (νt ))

τ

j=0

h=0

" H
ℓ
X

Hℓ
X

T #
∇θ log πθK (νt ) (a′j |s′j )

,

(58)

j=0

h=0

where τ ′ represents the trajectory for which the inner product term is maximum and thereby upper-bounding with that leads
to expression in equation (58) Next we upper-bound the norm of the ∥∇2ν,θ Vs (νt , θK (νt ))∥ as∥∇2ν,θ Vs (νt , θK (νt ))∥ ≤

Hℓ
X


T
Hℓ
X
γ h · ∇ν rνt (s′h , a′h ) · ∇θ log πθK (νt ) (a′j |s′j )
j=0

h=0

≤

Hℓ
X

γ h · ∇ν rνt (s′h , a′h ) ·HℓX∇θ log πθK (νt ) (a′j |s′j )

j=0

h=0

≤ Hℓ2 Lr B,(59)

where we apply the Cauchy-Schwarz inequality to get the equation in the second line. We apply triangle inequality with
Assumptions 2 and 3 to get the final bound for the mixed hessian term.
18Submission and Formatting Instructions for ICML 2023

11.2.2. P ROOF OF L EMMA 4.2 S TATEMENT ( II )Proof. We start by considering the term
∇2θ Vs (ν, θ∗ (ν)) − ∇2θ Vs (ν, θK (ν)) = EP (τ ;θ∗ (ν)) f1 (τ ) − EP (τ ;θK (ν)) f2 (τ ),(60)

where we define

f1 (τ ) =HℓXγ h−1 · rνt (sh , ah ) · 

f2 (τ ) =∇2θ log πθ∗ (ν) (aj |sj )

(61)

j=0

h=0
Hℓ
X

Hℓ
X


γ h−1 · rνt (sk , ak ) · HℓX∇2θ log πθK (ν) (aj |sj ) .

(62)

j=0

h=0

Subsequently, we write the norm of the equation (60) into 2 parts as∥∇2θ Vs (ν, θ∗ (ν)) − ∇2θ Vs (ν, θK (ν))∥= ∥EP (τ ;θ∗ (ν)) f1 (τ ) − EP (τ ;θK (ν)) f1 (τ ) + EP (τ ;θK (ν)) f1 (τ ) − EP (τ ;θK (ν)) f2 (τ )∥≤ ∥T1 ∥ + ∥T2 ∥,(63)

where, T1 = EP (τ ;θ∗ (ν)) f1 (τ ) − EP (τ ;θK (ν)) f1 (τ ) and T2 = EP (τ ;θK (ν)) f1 (τ ) − EP (τ ;θK (ν)) f2 (τ ). We next, upper-boundthe individual terms to get the final Lispchitz constant.
First, we focus on the first term of inequality i.e T1 given as
T1 = EP (τ ;θ∗ (ν)) f1 (τ ) − EP (τ ;θK (ν)) f1 (τ )(64)

≤ sup[EP (τ ;θ∗ (ν)) f1 (τ ) − EP (τ ;θK (ν)) f1 (τ )]f1≤ Lf1 χ1 dT V (P (τ ; θ∗ (ν)), P (τ ; θK (ν)))
Hℓ≤ Lf1 χ1
L2 ∥θ∗ (ν) − θK (ν)∥,
2where we convert the inequality first to a standard Integral Probability Metric form by taking the supremum and then dividing
and multiplying with the Lipschitz constant Lf1 from equation (78), we get the final expression in terms of Total variation
where χ1 is the constant. Then, we upper-bounded the total variation using the results from equation (57) to get the final
expression.
Now, we proceed to the second term of the equation T2 and derive an upper bound as
T2 =X

P (τ ; θK (ν))(f1 (τ ) − f2 (τ ))

(65)

τ

≤ f1 (τ ′ ) − f2 (τ ′ )
=

Hℓ
X
h=0



Hℓ
X
γ h−1 · rνt (sh , ah ) ·  (∇2θ log πθ∗ (ν) (aj |sj ) − ∇2θ log πθK (ν) (aj |sj ))
j=0

,
19

Submission and Formatting Instructions for ICML 2023

where we consider the trajectory τ ′ in the sum with the maximum value and upper bound by that to get equation (65).

∥T2 ∥ ≤ ∥

Hℓ
X



Hℓ
X
γ h−1 · rνt (sh , ah ) ·  (∇2θ log πθ∗ (ν) (aj |sj ) − ∇2θ log πθK (ν) (aj |sj )) ∥

≤∥

Hℓ
X

(66)

j=0

h=0



Hℓ
X
γ h−1 · rνt (sh , ah ) ·  (∇2θ log πθ∗ (ν) (aj |sj ) − ∇2θ log πθK (ν) (aj |sj )) ∥
j=0

h=0



Hℓ
Hℓ
X
X
≤
γ h−1 · ∥rνt (sh , ah )∥  (∥∇2θ log πθ∗ (ν) (aj |sj ) − ∇2θ log πθK (ν) (aj |sj )∥)
j=0

h=0

≤

Hℓ
X

γ h−1 · ∥rνt (sh , ah )∥Hℓ L2 ∥θ∗ (ν) − θK (ν)∥

h=0

≤ L2 RHℓ ∥θ∗ (ν) − θK (ν)∥

Hℓ
X

γ h−1

h=0

= L2 RHℓ2 ∥θ∗ (ν) − θK (ν)∥,
where we use Cauchy-Schwartz and triangle inequality repetitively to get to the third inequality. Next, we use Assumption 3
on the Lipschitzness of the gradient of the score function and the bounded reward norm max(s,a) ∥rν (s, a)∥ = R to get thenext inequality. Finally, we use the upper bound on the geometric series to obtain the final expression.
Adding equations (64) and (66), we get the∥∇2θ Vs (ν, θ∗ (ν)) − ∇2θ Vs (ν, θK (ν))∥ ≤ L′ ∥θ∗ (ν) − θK (ν)∥(67)

where, L′ = Lf1 χ1 H2ℓ L2 + L2 RHℓ2 and Lf1 = L2 Hℓ2 R.
11.2.3. P ROOF OF L EMMA 4.2 S TATEMENT ( III )Proof. We start by considering the term
∇2ν,θ Vs (ν, θ∗ (ν)) − ∇2ν,θ Vs (ν, θK (ν)) ≤ EP (τ ;θ∗ (ν)) f3 (τ ) − EP (τ ;θK (ν)) f4 (τ )(68)

where we define

f3 (τ ) =f4 (τ ) =

Hℓ
X


T
Hℓ
X
γ h−1 · ∇ν rν (sh , ah ) · ∇θ log πθ∗ (ν) (aj |sj )

h=0

j=0

Hℓ
X


T
Hℓ
X
γ h−1 · ∇ν rν (sh , ah ) · ∇θ log πθK (ν) (aj |sj ) .

h=0

j=0

(69)

(70)

Subsequently, we write the norm of the equation (68) into 2 parts as∥∇2ν,θ Vs (ν, θ∗ (ν)) − ∇2ν,θ Vs (ν, θK (ν))∥ ≤ ∥EP (τ ;θ∗ (ν)) f3 (τ ) − EP (τ ;θK (ν)) f3 (τ )(71)

+ EP (τ ;θK (ν)) f3 (τ ) − EP (τ ;θK (ν)) f4 (τ )∥≤ ∥T3 ∥ + ∥T4 ∥,where, T3 = EP (τ ;θ∗ (ν)) f3 (τ ) − EP (τ ;θK (ν)) f3 (τ ) and T4 = EP (τ ;θK (ν)) f3 (τ ) − EP (τ ;θK (ν)) f4 (τ ). We next, upper-boundthe individual terms to get the final Lispchitz constant.
20

Submission and Formatting Instructions for ICML 2023

First, we focus on the first term of inequality i.e T3 given as
T3 = EP (τ ;θ∗ (ν)) f3 (τ ) − EP (τ ;θK (ν)) f3 (τ )(72)

≤ sup[EP (τ ;θ∗ (ν)) f3 (τ ) − EP (τ ;θK (ν)) f1 (τ )]f3≤ Lf3 χ2 dT V (P (τ ; θ∗ (ν)), P (τ ; θK (ν)))
Hℓ
L2 ∥θ∗ (ν) − θK (ν)∥,≤ Lf3 χ22
where we convert the inequality first to a standard Integral Probability Metric form by taking the supremum and then dividing
and multiplying with the Lipschitz constant Lf3 from equation (81), and then we get the final expression in terms of Total
variation where χ2 is the constant. Then, we upper-bounded the total variation using the results from equation (57) to get
the final expression.
Now, we proceed to the second term of the equation T4 and derive an upper bound as
T4 =X

P (τ ; θK (ν))(f3 (τ ) − f4 (τ ))

(73)

τ

≤ f3 (τ ′ ) − f4 (τ ′ )
=

Hℓ
X

T
Hℓ
X
γ h−1 · ∇ν rν (sh , ah ) ·  (∇θ log πθ∗ (ν) (aj |sj ) − ∇θ log πθK (ν) (aj |sj ))


j=0

h=0

where we consider the trajectory τ ′ in the sum with the maximum value and upper bound by that to get equation (73).
Next, we upper-bound the norm ∥T4 ∥ as∥T4 ∥ ≤

Hℓ
X


T
Hℓ
X
γ h−1 · ∇ν rν (sh , ah ) ·  (∇θ log πθ∗ (ν) (aj |sj ) − ∇θ log πθK (ν) (aj |sj ))
j=0

h=0

≤

Hℓ
X


γ h−1 · ∥∇ν rν (sh , ah )∥ ≤∥∇θ log πθ∗ (ν) (aj |sj ) − ∇θ log πθK (ν) (aj |sj ))∥

j=0

h=0
Hℓ
X

Hℓ
X

γ h−1 · ∥∇ν rν (sh , ah )∥Hℓ L1 ∥θ∗ (ν) − θK (ν)∥

h=0

≤ L1 Lr Hℓ ∥θ∗ (ν) − θK (ν)∥

Hh
X

γ h−1

h=0

≤ L1 Lr Hℓ2 ∥θ∗ (ν) − θK (ν)∥,(74)

where we use Cauchy-Schwartz and triangle inequality repetitively to get to the third inequality. Next, we use Assumption
3 on the Lipschitzness of the gradient of the score function and the bounded reward R (cf. Assumption 2). Finally, we
upper-bound sum of this geometric series to obtain the final expression. Adding equations (72) and (74), we get the∥∇2ν,θ Vs (ν, θ∗ (νt )) − ∇2ν,θ Vs (νt , θK (ν))∥ ≤ L′′ ∥θ∗ (ν) − θK (ν)∥
where, L′′ = Lf3 χ2 H2ℓ L2 + L1 Lr Hℓ2 and Lf3 = Lr L1 Hℓ2 .
21

(75)

Submission and Formatting Instructions for ICML 2023

12. Additional Supporting Lemmas
12.1. Proof of Lispchitzness for f1 (·) defined in (61)Here, we prove that the function denoted as f1 (θ∗ (ν)) =Lispchitz continuous w.r.t θ with Lipschitz constant Lf1 i.e

PHℓ

h−1
· rνt (sh , ah ) ·
h=0 γ

P
Hℓ


2
∗ (ν) (aj |sj )
∇
log
π
is
θ
j=0 θ

∥f1 (θ∗ (ν)) − f1 (θK (ν))∥ ≤ Lf1 ∥θ∗ (ν) − θK (ν)∥.

(76)

First, we begin with the difference term as :f1 (θ∗ (ν)) − f1 (θK (ν)) =

Hℓ
X



Hℓ
X
γ h−1 rνt (sh , ah )  (∇2θ log πθ∗ (ν) (aj |sj ) − ∇2θ log πθK (ν) (aj |sj )) .

(77)

j=0

h=0

Subsequently, taking the norm we get∥f1 (θ∗ (ν)) − f1 (θK (ν))∥
=

Hℓ
X



Hℓ
X
γ h−1 rνt (sh , ah )  (∇2θ log πθ∗ (ν) (aj |sj ) − ∇2θ log πθK (ν) (aj |sj ))
j=0

h=0

≤

Hℓ
X



Hℓ
X
γ h−1 · ∥rνt (sh , ah )∥ · ∥∇2θ log πθ∗ (ν) (aj |sj ) − ∇2θ log πθK (ν) (aj |sj )∥
j=0

h=0

≤

Hℓ
X

γ h−1 · ∥rνt (sh , ah )∥L2 Hℓ ∥θ∗ (ν) − θK (ν)∥

h=0

≤ L2 Hℓ2 R∥θ∗ (ν) − θK (ν)∥,(78)

where we use Cauchy-Schwartz and triangle inequality repeatedly to get the subsequent inequalities. In the third inequality,
we use the Lipschitzness assumption of ∇2θ log πθ∗ (ν) (aj |sj ) from Assumption 3 and finally using the boundedness of the
reward values and upper-bounding the Geometric series, we get the final expression. Thus f1 (θ∗ (ν)) is Lispchitz continuous
w.r.t θ with Lipschitz constant Lf1 = L2 Hℓ2 R.
12.2. Proof of Lispchitzness for f3 (·) defined in (69)P
TPHℓ h−1HℓHere, we prove that the function denoted as f3 (θ∗ (ν)) = h=0γ· ∇ν rν (sh , ah ) ·is
j=0 ∇θ log πθ ∗ (ν) (aj |sj )Lispchitz continuous w.r.t θ with Lipschitz constant Lf3 i.e∥f3 (θ∗ (ν)) − f3 (θK (ν))∥ ≤ Lf3 ∥θ∗ (ν) − θK (ν)∥(79)

First, we begin with the difference term as :f3 (θ∗ (ν)) − f3 (θK (ν))
=

Hℓ
X
h=0


T
Hℓ
X
γ h−1 · ∇ν rν (sh , ah ) ·  (∇θ log πθ∗ (ν) (aj |sj ) − ∇θ log πθ∗ (ν) (aj |sj ))
j=0

22

(80)

Submission and Formatting Instructions for ICML 2023

Subsequently, taking the norm we get∥f3 (θ∗ (ν)) − f3 (θK (ν))∥
=∥

Hℓ
X


T
Hℓ
X
γ h−1 · ∇ν rν (sh , ah ) ·  (∇θ log πθ∗ (ν) (aj |sj ) − ∇θ log πθ∗ (ν) (aj |sj )) ∥
j=0

h=0



Hℓ
Hℓ
X
X
≤
γ h−1 · ∥∇ν rν (sh , ah )∥ ·  (∥∇θ log πθ∗ (ν) (aj |sj ) − ∇θ log πθ∗ (ν) (aj |sj )∥)
j=0

h=0

≤ Lr L1 Hℓ ∥θ∗ (ν) − θK (ν)∥

Hℓ
X

γ h−1

h=0
2 ∗
K
≤ Lr L1 Hℓ ∥θ (ν) − θ (ν)∥,(81)

where we use Cauchy-Schwartz and triangle inequality repeatedly to get the subsequent inequalities. In the third inequality,
we use the Lipschitzness assumption of ∇2θ log πθ∗ (ν) (aj |sj ) from Assumption 3. Finally, using the Assumption 2 i.e reward
Lipschitzness and upper-bounding the Geometric series, we get the final expression. Thus f3 (θ∗ (ν)) is Lispchitz continuous
w.r.t θ with Lipschitz constant Lf3 = Lr L1 Hℓ2 .
12.3. Proof of Smoothness Condition on the Value function
Here, we prove the smoothness of the value function i.e. the gradient of the value function is Lispchitz continuous w.r.t θ
with Lipschitz constant L1 . We begin with the definition of the gradient difference of the value function from the equation(27) as:∇θ Vs (ν, θ∗ (ν)) − ∇θ Vs (νt , θK (ν)) = EP (τ ;θ∗ (ν)) V1 (τ ) − EP (τ ;θK (ν)) V2 (τ )(82)

= EP (τ ;θ∗ (ν)) V1 (τ ) − EP (τ ;θK (ν)) V1 (τ )+ EP (τ ;θK (ν)) [V1 (τ ) − V2 (τ )]= Σ1 + Σ2P
PHℓ −1 h−1
Hℓ −1∗
where, first we substitute V1
=
γ
r
(s
,
a
)
∇
log
π
(a
|s
)
and V2=θνt h hθ (νt ) j j
h=0
j=0
P

PHℓ −1 h−1
Hℓ −1
rνt (sh , ah )h=0 γ
j=0 ∇θ log πθ K (νt ) (aj |sj ) . Subsequently, by adding and subtracting EP (τ ;θ K (ν)) V2 (τ ), we
get the final expression, where Σ1 = EP (τ ;θ∗ (ν)) V1 (τ ) − EP (τ ;θK (ν)) V1 (τ ) and Σ2 = EP (τ ;θK (ν)) [V1 (τ ) − V2 (τ )]. Now,
first, we derive the Lipschitz constant for V1 as∥V1 (θ∗ (ν) − V1 (θK (ν))∥
≤

H
ℓ −1
X


γ h−1 ∥rνt (sh , ah )∥ 

H
ℓ −1
X


∥∇θ log πθ∗ (νt ) (aj |sj ) − ∇θ log πθK (νt ) (aj |sj )∥

j=0

h=0

≤ Hl2 RL1 ∥θ∗ (ν) − θK (ν)∥,(83)

where we first use Cauchy-Schwartz and triangle inequality to get the first inequality. Next, we upper-bound the rewardwith R from Assumption 2, Lipschitzness of policy gradient from Assumption 3, and finally upper-bounding the Geometricseries, we get the final expression. The Lipschitz constant L5 = Hl2 RL1 .
We can subsequently upper-bound Σ1 with the total variation distance asΣ1 ≤ sup[EP (τ ;θ∗ (ν)) V (τ ) − EP (τ ;θK (ν)) V (τ )]V≤ L5 dT V (P (τ ; θK (ν)), P (τ ; θ∗ (ν)))
HL2 ∗
≤ L5
∥θ (ν) − θK (ν)∥
2
23

(84)

Submission and Formatting Instructions for ICML 2023

where first we divide by the Lipschitz constant of the function and subsequently upper-bound with the Total Variation
distance. Finally, we substitute the total variation distance expression from the equation (57) to get the final expression.
Now, the second term can be written asΣ2 = EP (τ ;θK (ν)) [V1 (τ ) − V2 (τ )] ≤ V1 (τ ′ ) − V2 (τ ′ ),(85)

where we take the trajectory with the maximum difference and upper-bound the term. Subsequently, ∥Σ2 ∥ ≤ ∥V1 (τ ′ ) −V2 (τ ′ )∥ = Hl2 RLθ1 ∥θ∗ (ν) − θK (ν)∥ from equation (83).
Finally, the norm of the gradient difference of the value function from equation (82) as∥∇θ Vs (ν, θ∗ (ν)) − ∇θ Vs (νt , θK (ν))∥ ≤ L6 ∥θ∗ (ν) − θK (ν)∥

(86)

2
2
where L6 = L5 HL
2 + L5 and L5 = Hl RLθ1

12.4. Upper-bound on the Norm of the hessian defined in (28)Here, we prove an upper-bound on the norm of the hessian defined in (28) given as

#
" H
Hℓ
ℓ
X
X
∇2θ Vs (νt , θK (νt )) = EP (τ ;θK (νt ))γ h−1 rνt (sh , ah ) ∇2θ log πθK (νt ) (aj |sj )

≤

Hℓ
X


γ h−1 rνt (sh , ah ) HℓX∇2θ log πθK (νt ) (aj |sj )

X

P (τ ; θK (νt ))

τ

j=0

h=0

(87)

j=0

h=0



Hℓ
Hℓ
X
X
=
γ h−1 rνt (sh , ah ) ∇2θ log πθK (νt ) (aj |sj )
j=0

h=0

where, first we upper-bound the function with the trajectory which has the maximum inner-value. Next we, upper-bound thenorm as


Hℓ
Hℓ
X
X
∥∇2θ Vs (νt , θK (νt ))∥ ≤ ∥γ h−1 rνt (sh , ah ) ∇2θ log πθK (νt ) (aj |sj ) ∥
(88)
j=0

h=0

≤

Hℓ
X


γ h−1 ∥rνt (sh , ah )∥ HℓX

∥∇2θ log πθK (νt ) (aj |sj )∥

j=0

h=0

≤ Hl2 RL1π
where, we first upper-bound with successive application of Cauchy-Schwartz and Triangle inequality to get the second
inequality. Finally, with the upper bound on ∥rνt (sh , ah )∥ ≤ R from Assumption (15), Lipschitzness of policy gradientsfrom Assumption 2 and upper-bounding the geometric series, we get the final expression.
12.5. Upper-bound on the Norm of the hessian defined in (28)Here, we prove an upper-bound on the norm of the mixed second-order Jaccobian term defined in (29) given as
"H
X
#
Hℓ
ℓ
X
∇2ν,θ Vs (νt , θK (νt )) = EP (τ ;θK (νt ))γ h−1 ∇ν rνt (sh , ah )[∇θ log πθK (νt ) (aj |sj )]T
j=0

h=0

X

Hℓ
Hℓ
X
h−1
T
≤
γ
∇ν rνt (sh , ah )[∇θ log πθK (νt ) (aj |sj )]
h=0

j=0

24

(89)

Submission and Formatting Instructions for ICML 2023

where first we upper-bound the function with the trajectory which has the maximum inner value. Next, we, upper-bound thenorm as
X

Hℓ
Hℓ
X
2
K
h−1
T
∥∇ν,θ Vs (νt , θ (νt ))∥ ≤ ∥ ≤γ∇ν rνt (sh , ah )[∇θ log πθK (νt ) (aj |sj )] ∥
(90)
j=0

h=0
Hℓ
X

X

Hℓ
h−1
γ
∥∇ν rνt (sh , ah )∥∥∇θ log πθK (νt ) (aj |sj )∥ ∥

≤

j=0

h=0

≤ Hl2 Lr B
where we first upper-bound with the successive application of Cauchy-Schwartz and Triangle inequality to get the second
inequality. Finally, with the upper bound on ∥∇ν rνt (sh , ah )∥ ≤ Lr from Lipschitzness Assumption (15), Lipschitzness ofpolicy parametrization from Assumption 2 and upper-bounding the geometric series, we get the final expression.
12.6. Proof of Lemma 4.3
Proof. Let us start by first deriving the upper bounds for the terms ϕ1 and ϕ2 as defined in the equation (41) as follows. For
ϕ1 (τ ), we have∥ϕ1 (τ )∥ = U (τ ) ·H−1
X[−∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θ∗ (ν))−1 ∇θ fh (θ∗ (ν))] .

(91)

h=0

We define the term κ which explains the relative conditioning of the two matrix norms. We define the mixed condition
number as∥∇2v,θ Vs (ν, θ∗ (ν))∥
Hl Lr B
κ=
(92)
≤
2
∗
∥∇θ Vs (ν, θ (ν))∥
lπ (1 − γ)
Next to upper-bound the second term relating to the difference in ∥ϕ1 − ϕ2 ∥ in the equation, we proceed first by upperbounding the product difference∥∆∥ = ∥∆1 − ∆2 ∥(93)

∆1 =∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θ∗ (ν))−1 ∇θ fh (θ∗ (ν)),(94)∆2 =∇2v,θ Vs (ν, θK (ν))∇2θ Vs (ν, θK (ν))−1 ∇θ fh (θK (ν)).

(95)

where we define

Also, for simplicity of notations lets take ψ1=∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θ∗ (ν))−1
2
K
2
K
−1
∇v,θ Vs (ν, θ (ν))∇θ Vs (ν, θ (ν)) , which thus (93) boils down to upper-boundingand ψ2∆ = ∥ψ1 fh (θ∗ (ν)) − ψ2 fh (θK (ν))∥
∗

K

=
(96)

K

K

= ∥ψ1 fh (θ (ν)) − ψ1 fh (θ (ν)) + ψ1 fh (θ (ν)) − ψ2 fh (θ (ν))∥≤ ∥ψ1 ∥∥fh (θ∗ (ν)) − fh (θK (ν))∥ + ∥ψ1 − ψ2 ∥∥fh (θK (ν))∥≤ κL1 ∥θ∗ (ν) − θK (ν)∥ + L2 ∥ψ1 − ψ2 ∥,where, first we expand add and subtract the term ψ1 fh (θK (ν)), and subsequently by applying Cauchy-Schwartz and triangle
inequality, we get to the third inequality. For the final inequality, we apply equation and Lispchitzness Assumptions (15) to
get the final expression in equation (96). Next, we focus on upper bounding the second term of the expression specifically∥ψ1 − ψ2 ∥∥ψ1 − ψ2 ∥ =∥∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θ∗ (ν))−1 − ∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θK (ν))−1+ ∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θK (ν))−1 − ∇2ν,θ Vs (ν, θK (ν))∇2θ Vs (ν, θK (ν))−1 ∥=∥∇2ν,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θ∗ (ν))−1 − ∇2v,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θK (ν))−1 ∥+ ∥∇2ν,θ Vs (ν, θ∗ (ν))∇2θ Vs (ν, θK (ν))−1 − ∇2v,θ Vs (ν, θK (ν))∇2θ Vs (ν, θK (ν))−1 ∥=Ψ21 + Ψ22 ,
25

(97)
(98)

Submission and Formatting Instructions for ICML 2023

where we expand the definition of ∥ψ1 − ψ2 ∥ and subsequently apply triangle inequality and Cauchy-Schwartz which thenboils to upper-bounding the sum of two terms Ψ21 + Ψ22 . For Ψ21 , we have
Ψ21 ≤ ∥∇2ν,θ Vs (ν, θ∗ (ν))∥∥∇2θ Vs (ν, θ∗ (ν))−1 − ∇2v,θ Vs (θK (ν))−1 ∥(99)≤ Lν,θ ∥∇2θ Vs (ν, θ∗ (ν))−1 − ∇2θ Vs (ν, θK (ν))−1 ∥= Lν,θ ∥∇2θ Vs (ν, θ∗ (ν))−1 (∇2θ Vs (ν, θ∗ (ν)) − ∇2θ Vs (ν, θK (ν)))∇2θ Vs (ν, θK (ν))−1 ∥≤ Lν,θ ∥∇2θ Vs (ν, θ∗ (ν))∥−1 ∥(∇2θ Vs (ν, θ∗ (ν)) − ∇2θ Vs (ν, θK (ν)))∥∥∇2θ Vs (ν, θK (ν))∥−1≤

Lν,θ L′ ∗∥θ (ν) − θK (ν)∥
lπ2where we use Cauchy-Schwartz inequality and triangle inequality iteratively to get the final inequality. Next, we use theupper bounds and lower bounds of the hessian and mixed hessian matrices defined in Assumptions to get the final expression.
Now, Lν,θ = Hl2 Lr B from equation (90), L′ = Lf1 χ1 H2ℓ L2 + L2 RHℓ2 and Lf1 = L2 Hℓ2 R from equation (67). Finally,
the second-term Ψ22 from equation (97) can be upper-bounded asΨ22 ≤ ∥∇2ν,θ Vs (ν, θ∗ (ν) − ∇2v,θ Vs (ν, θK (ν))∥∥∇2θ Vs (ν, θK (ν))∥−1≤L′′ ∗∥θ (ν) − θK (ν)∥,
lπ

(100)

where, similarly we use triangle inequality with Cauchy-Schwartz to get the final upper-bound of equation (100). Now,
L′′ = Lf3 χ2 H2ℓ L2 + L1 Lr Hℓ2 and Lf3 = Lr L1 Hℓ2 from equation (75)Now, combining equations (99) and (100), we get the final upper-bound of the ∥ψ1 − ψ2 ∥ in equation (97) as∥ψ1 − ψ2 ∥ ≤ (L′′Lν,θ L′+)∥θ∗ (ν) − θK (ν)∥.
2
lπ
lπ

(101)

Hence, finally replacing the upper-bound of Ψ2 from equation (101) in equation (96) to obtain the upper-bound on thefunction difference term ∆ as∆ ≤ κL1 ∥θ∗ (ν) − θK (ν)∥ + (Lν,θ L′L′′+)∥θ∗ (ν) − θK (ν)∥
lπ2
lπ

(102)

= γ1 ∥θ∗ (ν) − θK (ν)∥,
L

L′

′′

with γ1 := κL1 + ν,θ
+ Llπ Hence, with the above bounds, we proceed to upper-bound the term II in equation (39) i.e2lπ∥Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ ) − ϕ2 (τ )]∥ as∥Eτ ∼P (τ ;θK (ν)) [ϕ1 (τ ) − ϕ2 (τ )]∥ ≤ ∥ϕ1 (τ ′ ) − ϕ1 (τ ′ )∥= ∥U (τ ) ·

(103)

H−1
X

(∆1 − ∆2 )∥

h=0

≤∥

H−1
X
h′ =0
2

u(sh , ah )∥∥H−1
X(∆1 − ∆2 )∥h=0≤ H ũ∥∆1 − ∆2 ∥≤ H 2 ũγ1 ∥θ∗ (ν) − θK (ν)∥,
where first we select the trajectory τ ′ with the maximum sum and subsequently using the Cauchy-Schwartz inequality we
get the second equation. Based on the assumption of bounded utility u(s, a) ≤ ũ, ∀(s, a), we get the third equation andfinally using the upper bound of ∆1 − ∆2 from equation (102), we get the final expression for equation (103).
26

Submission and Formatting Instructions for ICML 2023

12.7. Proof of Lemma 4.4
Proof. Here, we derive an upper bound on the tracking term due to the use of surrogate gradients ∥θ∗ (νt ) − θK (νt )∥. To
begin the proof, we start with the smoothness in the value function shown in equation (86), as
Vs (νt , θk+1 (νt )) ≤ Vs (νt , θk (νt )) + ⟨∇θ Vs (νt , θk+1 (νt )), θk+1 (νt ) − θk (νt )⟩L6 k+1+∥θ(νt ) − θk (νt )∥2 .
2
where L6 = L5

(104)

Hℓ2 L2+ L5 and L5 = Hl RLθ1 . Now, from the update of θ from Algorithm 2, we know that2

θk+1 (νt ) = θk (νt ) − αℓ ∇θ Vs (νt , θk (νt )).

(105)

Replacing the update in equation (104), we haveVs (νt , θk+1 (νt )) ≤Vs (νt , θk (νt )) + ⟨∇θ Vs (νt , θk+1 (νt )), −αℓ ∇θ Vs (νt , θk (νt ))⟩L6∥ − αℓ ∇θ Vs (νt , θk (νt ))∥2
+
2


αℓ L6
k
∥∇θ Vs (νt , θk (νt ))∥2 ,=Vs (νt , θ (νt )) − αℓ 1 −
2

(106)

where we expand the expression after replacing the update in equation (104).
Next, from Assumption 4, we note that for the value function, it holds that∥∇θ Vs (ν, θk (ν))∥2 ≥µ(Vs (ν, θk (ν)) − Vs (ν, θ∗ (ν))).
2

(107)

Assumption 4 ensures that the objective satisfies the gradient dominance or the PL condition, which can be satisfied inpractice for various settings. For instance, Assumption 4 is satisfied in our setting for softmax policy parametrization as
detailed in (Mei et al., 2020b, Lemma 8). Now, replacing the PL condition in equation (106), we haveαℓ L6 µ) (Vs (ν, θk (ν)) − Vs (ν, θ∗ (ν)))
22= −α3 (Vs (νt , θk (νt )) − Vs (νt , θ∗ (νt ))),Vs (νt , θk+1 (νt )) − Vs (νt , θk (νt )) ≤ −αℓ (1 −(108)

where, after replacing the PL condition in equation (106), we substitute α3 = αℓ (1 − αℓ2L6 ) µ2 for simplicity of calculations.
Vs (νt , θk+1 (νt )) − Vs (νt , θ∗ (νt ))) ≤ (1 − α3 )(Vs (νt , θk (νt )) − Vs (νt , θ∗ (νt )))Vs (νt , θK (νt )) − Vs (νt , θ∗ (νt ))) ≤ (1 − α3 )K (Vs (νt , θ0 (νt )) − Vs (νt , θ∗ (νt ))),

(109)

where the first equation comes from algebraic manipulation and applying the equation recursively, we get the second
inequality, assuming 0 ≤ α3 ≤ 1. Now, we note that from the smoothness of value function, we have the upper boundVs (νt , θ0 (νt )) − Vs (νt , θ∗ (νt )) ≤L6 ∗∥θ (νt ) − θ0 (νt )∥2 ,
2

(110)

2
2
where L6 = L5 HL
2 + L5 , L5 = Hl RL1 and we use the Lipschitz smoothness assumption and expand along the point∗
νt , θ , for which the gradient term vanishes. Also, since PL implies quadratic growth, it holds thatVs (νt , θK (νt )) − Vs (νt , θ∗ (νt )) ≥µ K
∥θ (νt ) − θ∗ (νt )∥2 .
2

(111)

L6
Z,
µ

(112)

Now, substituting the equations (110), (111) in (109) to obtain∥θK (νt ) − θ∗ (ν)∥2 ≤ (1 − α3 )K

where, Z := maxν ∥θ0 − θ∗ (ν)∥2 , α3 = αℓ (1 − αℓ2L6 ) µ2 and L6 = L5 Hℓ2L2 + L5 and L5 = Hl2 RL1 .
27

Submission and Formatting Instructions for ICML 2023

13. Discuss of Assumption 3 and Assumption 4• Assumption 3 ensure certain properties of the policy parametrization such as Lipschitz policy, bounded score function,
Lipschitz score function, and Lipschitz Hessian of the log of the policy. We remark that these assumptions are not
restrictive and satisfied in practice for practical classes of policies. For example, this assumption is satisfied for softmax
policy parametrization.
exp θT ϕ(s, a)T′a′ exp θ ϕ(s, a )πθ (a|s) = P(113)

where, first we write down the expression of the softmax policy gradient parametrization with function approximation
ϕ. Subsequently, taking log and taking the gradient w.r.t to the policy parameterization, we get
X1
exp θT ϕ(s, a′ )ϕ(s, a′ )T′
a′ exp θ ϕ(s, a ) a′Xπθ (a|s)ϕ(s, a′ )= ϕ(s, a) −∇θ log πθ (a|s) = ϕ(s, a) − P

(114)

a′

P
where we can denote ϕbs,a′ = a′ πθ (a|s)ϕ(s, a′ ). Now, taking the norm on the LHS of equation (114) we get∥∇θ log πθ (a|s)∥ ≤ ∥ϕ(s, a)∥ + ∥ϕ(s, a′ )∥(115)

where, we apply triangle inequality to get the final bound, which imposes certain constraints on the norm of the function
approximation ∥ϕ(s, a)∥ ≤ ζ1 which is a common assumption in various scenarios.
∥∇θ πθ (a|s)∥ ≤ kζ1

(116)

1
where, we expand ∇θ log πθ (a|s) = πθ (a|s)∇θ πθ (a|s)• Assumption 4 ensures that the objective function satisfies some geometric properties such as PL condition. We remarkthat the value function satisfies PL condition with softmax policy parametrization (see (Mei et al., 2020b, Lemma 8)).
Further, a property that we need for our analysis to hold is that the Hessian of the objective function has all non-zero
eigenvalues. This assumption holds for softmax parametrizations. Now, in order to show that we first consider thesoftmax-parametrization considered in equation (114), we first compute the hessian asX∇2θ log πθ (a|s) = ∇θ [ϕ(s, a) −πθ (a|s)ϕ(s, a′ )]
(117)
a′

=−

X

∇θ πθ (a|s)ϕ(s, a′ )Ta′= −Eπ [∇θ log πθ (a|s)ϕ(s, a′ )T ]b a′ )ϕ(s, a′ )T − ϕ(s, a)ϕ(s, a′ )T ]= Eπ [ϕ(s,Now, finally, we substitute this to the equation of hessian of the value function in equation (117)

#
"
H
ℓ −1
X
∇2θ Vs (νt , θK (νt )) = EP (τ ;θ) R(τ ) ∇2θ log πθK (νt ) (aj |sj )

(118)

j=0

"



= EP (τ ;θ) Eπ R(τ ) 

H
ℓ −1
X

#
b a′ )ϕ(s, a′ )T − ϕ(s, a)ϕ(s, a′ )T ]
[ϕ(s,

j=0

From the above, it is evident to ensure non-singular eigenvalues, we need to ensure non-singularity for the function
approximation matrix ϕϕT which has been a standard assumption in several settings (Sutton et al., 2009; Maei et al.,2009)28