Proximal Point Imitation Learning
Abstract
This work develops new algorithms with rigorous efficiency guarantees for infinite horizon imitation learning (IL) with linear function approximation without restrictive coherence assumptions.
We begin with the minimax formulation of the problem and then outline how to leverage classical tools from optimization, in particular, the proximal-point method (PPM) and dual smoothing, for online and offline IL, respectively.
Thanks to PPM, we avoid nested policy evaluation and cost updates for online IL appearing in the prior literature.
In particular, we do away with the conventional alternating updates by the optimization of a single convex and smooth objective over both cost and Q -functions.
When solved inexactly, we relate the optimization errors to the suboptimality of the recovered policy.
As an added bonus, by re-interpreting PPM as dual smoothing with the expert policy as a center point, we also obtain an offline IL algorithm enjoying theoretical guarantees in terms of required expert trajectories.
Finally, we achieve convincing empirical performance for both linear and neural network function approximation.
Introduction
This work is concerned with the prototypical setting of imitation learning (IL) where
An expert provides demonstrations of state-action pairs in an environment.
The expert could be optimal or suboptimal with respect to an unknown cost/reward function.
The learner chooses distance measure between its policy to be learned and the expert empirical distribution estimated from demonstrations.
The learner employs an algorithm, which additionally may or may not use interactions with the environment, to minimize the chosen distance.
In IL, the central goal of the learner is to recover a policy competitive with expert with respect to the underlying unknown cost function.
IL is important for several real world applications like driving [62], robotics [88], and economics/finance [27] at the expense of following resources: (R1) expert demonstrations, (R2)
(optional) interactions with the environment where the expert collected the demonstrations, and (R3) computational resources for solving the problem template.
Interestingly, while there is a vast amount of literature using optimization ideas on the IL problem template, i.e. Lagrangian duality [51, 38, 59, 63, 64], resource guarantees are still widely missing since the optimization literature focuses on the resource (R3) where IL literature mainly focuses on the first two resources (R1) and (R2).
Our work leverages deeper connections between optimization tools and IL by showing how classical optimization tools can be applied in a linear programming formulation of IL problem guaranteeing efficiency in all (R1), (R2), (R3).
Our contributions: This work aims at designing an algorithm enjoying both theoretical guarantees and convincing empirical performance.
Our methodology is rooted in classical optimization tools and the LP approach to MDPs.
More precisely, the method uses the recently repopularized overparameterization technique to obtain the Q-function as a Lagrangian multiplier [77, 14] and solves the associated program using a PPM update with appropriately chosen Bregman divergences.
This results to an actor-critic algorithm, with the key feature that the policy evaluation step involves optimization of a single concave and smooth objective over both cost and Q -functions.
In this way, we avoid instability or poor convergence due to adversarial training [51, 122, 70, 105], and can also recover an explicit cost along with Q-function.
We further account for potential optimization errors, presenting an error propagation analysis that leads to rigorous guarantees for both online and offline setting.
For the context of linear MDPs [14, 121, 55, 22, 116, 7, 84], we provide explicit convergence rates and error bounds for the suboptimality of the learned policy, under mild assumptions, significantly weaker than those found in the literature until now.
To our knowledge, such guarantees in this setting are provided for the first time.
Finally, we demonstrate that our approach achieves convincing empirical performance for both linear and neural network function approximation.
Related Literature.
The first algorithm addressing the imitation learning problem is behavioral cloning [93].
Due to the covariate shift problem [98, 99], it has low efficiency in terms of expert trajectories (R1).
To address this issue, [100, 87, 4, 95, 111, 85, 123, 5, 68, 69] proposed to cast the problem as inverse reinforcement learning (IRL).
IRL improves the efficiency in terms of expert trajectories, at the cost of introducing the need of running reinforcement learning (RL) repetitively, which can be prohibitive in terms of environment samples (R2) and computation (R3).
A successive line of work started with [112] highlights that repeated calls to an RL routine can be avoided.
This work inspired generative adversarial imitation learning (GAIL) [51] and other follow-up works [38, 59, 63, 64] that leveraged optimization tools like primal-dual algorithms but did not try to deepen the optimization connections to derive efficiency guarantees in terms of all (R1),(R2),(R3).
Finally, a recent line of work [40, 57] in IL bypasses the need of optimizing over cost functions and thus avoids instability due to adversarial training.
Although these algorithms achieve impressive empirical performance in challenging high dimensional benchmark tasks, they are hampered by limited theoretical understanding.
This is the fundamental difference from our work, which enjoys both favorable practical performance and strong theoretical guarantees.
Existing model-free IL theoretical papers with global convergence guarantees assume either a finite horizon episodic
MDP setting [70], or tabular MDPs [105], or the infinite horizon case but with restrictive assumptions, such as linear quadratic regulator setting [21], continuous kernelized nonlinear regulator [26, 56], access to a generative model and coherence assumption on the choice of features [58, 14], bounded strong concentrability coefficients [122] or a linear transition law that can be completely specified by a finite-dimensional matrix [70].
On the other hand, we provide convergence guarantees and error bounds for the context of linear MDPs [14, 121, 55, 22, 116, 7, 84] under a mild feature excitation condition assumption.
Despite being linear, the transition law can still have infinite degrees of freedom.
To our knowledge, such guarantees in this setting are provided for the first time.
Our work applies the technique known as regularization in the online learning literature [6, 103] and Bregman proximal-point or smoothing in optimization literature [97, 82] to the LP formulation for MDPs [73, 35, 36, 17, 48, 49, 33, 34, 102, 91, 92, 1, 65, 30, 79, 115, 67, 13, 31, 55, 106].
From this perspective, we can see Deep Inverse Q-Learning [57] and IQ-Learn [40] that consider entropy regularization in the objective as smoothing using uniform distribution as center point.
In our case, we instead use as center point the previous iteration of the algorithm (for the online case) or the expert (for the offline case).
From the technical point of view, the most important related works are the analysis of REPS/QREPS [90, 14, 89] and O-REPS [124] that first pointed out the connection between REPS and
PPM.
We build on their techniques with some important differences.
In particular, while in the
LP formulation of RL, PPM and mirror descent [15, 47] are equivalent, recognizing that they are not equivalent in IL is critical for stronger empirical performance.
As an independent interest, our techniques can be used to improve upon the best rate for REPS in the tabular setting [89] and to extend the guarantees to linear MDPs.
In order to discuss in more detail our research questions and situate them among prior related theoretical and practical works, we provide in Appendix A an extended literature review.
Background 2.1 Markov Decision Processes
The RL environment and its underlying dynamics are typically abstracted as an MDP given by a tuple ( S , A , P, ν 0 , c , γ ) , where S is the state space, A is the action space, P : S×A → ∆ S is the transition law, ν 0 ∈ ∆ S is the initial state distribution, c ∈ [0 , 1] |S||A| is the cost, and γ ∈ (0 , 1) is the discount factor.
For simplicity, we focus on problems where S and A are finite but too large to be enumerated.
A stationary Markov policy π : S → ∆ A interacts with the environment iteratively, starting with an initial state s 0 ∼ ν 0 .
At round t , if the system is at state s t , an action a t ∼ π ( ·| s t ) is sampled and applied to the environment.
Then a cost c ( s, a ) is incurred, and the system transitions to the next state s t
+1 ∼ P ( ·| s, a ) .
The goal of RL is to solve the optimal control problem ρ ⋆ c ≜ min π ρ c ( π ) , where ρ c ( π ) ≜ (1 − γ ) ⟨ ν 0 , V π c ⟩ is the normalized total discounted expected cost of π .
The state value function V π c ∈
R |S| of π , given cost c , is defined by V π c ( s ) ≜
E πs (cid:104) (cid:80) ∞ t =
γ t c ( s t , a t ) (cid:105) , where E πs denotes the expectation with respect to the trajectories generated by π starting from s 0 = s .
The optimal value function V ⋆ c ∈ R |S| is defined by
V ⋆ c ( s ) ≜ min π V π c ( s ) .
The optimal state-action value function Q ⋆ c ∈ R |S||A| , given by Q ⋆ c ( s, a ) ≜ c ( s, a ) + γ (cid:80) s ′
V ⋆ c ( s ′ )
P ( s ′ | s, a ) , is known to characterize optimal behaviors.
Indeed V ⋆ c is the unique solution to the Bellman optimality equation V ⋆ c ( s ) =
min a Q ⋆ c ( s, a ) .
In addition, any deterministic policy π ⋆ c ( s ) = arg min a Q ⋆ c ( s, a ) is known to be optimal.
For every policy π , we define the normalized state-action occupancy measure µ π ∈ ∆ S×A , by µ π ( s, a ) ≜ (1 − γ ) (cid:80) ∞ t =
γ t P π ν
[ s t
=
s, a t = a ] , where P π ν 0 [ · ] denotes the probability of an event when following π starting from s 0 ∼ ν 0 .
The occupancy measure can be interpreted as the discounted visitation frequency of state-action pairs.
This allows us to write ρ c ( π ) = ⟨ µ π , c ⟩ .
2.2 Imitation Learning
Similarly to RL, the IL problem is posed in the MDP formalism, with the critical difference that the true cost c true is unknown.
Instead, we have access to a finite set of truncated trajectories sampled i.i.d.
by executing an expert policy π E in the environment.
The goal is to learn a policy that performs better than π E with respect to the unknown c true .
To this end, we adopt the apprenticeship learning formalism [4, 112, 50, 51, 105], which carries the assumption that c true belongs to a class of cost functions C .
We then seek an apprentice policy π A that outperforms the expert across C by solving the following optimization problem
[Equation]
(1) where d C ( π, π E )
[Equation]
π
[Equation]
) (cid:1) defines the C -distance between π and π E [51, 28, 122, 70].
Then, π A satisfies the goal of IL, since it holds that ρ c true ( π A ) − ρ c true ( π E ) ≤ ζ ⋆ ≤ 0 .
Intuitively, the cost class C distinguishes the expert from other policies.
The maximization in (1) assigns high total cost to non-expert policies and low total cost to π E [51], while the minimization aims to find the policy that matches the expert as close as possible with respect to d C .
By writing d C in its dual form ¯ d C ( µ π , µ π E ) ≜ max c ∈C (cid:0) ⟨ µ π , c ⟩ − (cid:10) µ π E , c (cid:11) (cid:1) , it can be interpreted as an integral probability metric [80, 60] between the occupancy measures µ π and µ π E .
Depending on how C is chosen, d C turns to a different metric of probability measures like the 1 -Wasserstein distance [117, 32] for C
= Lip 1 ( S ×A ) , the total variation for C = { c | ∥ c ∥ ∞ ≤ 1 } , or the maximum mean discrepancy for C
= { c | ∥ c
∥ H ≤ 1 } , where Lip 1 ( S × A ) denotes the space of 1 -Lipschitz functions on S × A , and ∥·∥ H denotes the norm of a reproducing kernel Hilbert space H [104].
In our theoretical analysis, we focus on linearly parameterized cost classes [111, 112, 51, 70, 105] of the form C ≜ { c w ≜ (cid:80) mi
=1 w i ϕ i | w ∈ W} , where { ϕ i } mi =1 ⊂ R |S||A| + are fixed feature vectors, such that ∥ ϕ i ∥ 1 ≤ 1 for all i ∈ [ m ] , and W is a a convex constraint set for the cost weights w .
This assumption is not necessarily restrictive as usually in practice the true cost depends on just a few key properties, but the desirable weighting that specifies how different desiderata should be traded-off is unknown [4].
Moreover, the cost features can be complex nonlinear functions that can be obtained via unsupervised learning from raw state observations [20, 29].
The matrix Φ ≜ [ ϕ. . .
ϕ m ] gives rise a feature expectation vector (FEV) ρ Φ ( π ) ≜ ( ρ ϕ 1 ( π E ) , . . .
, ρ ϕ m ( π E ))
T ∈ R m of a policy π .
Then, by choosing W to be the ℓ 2 unit ball B m 1 ≜ { w ∈ R m | ∥ w ∥ 2 ≤ 1 } [4], we get a feature expectation matching objective d C ( π, π π E ) =
∥ ρ Φ ( π ) −
ρ Φ ( π E ) ∥ 2 , while for W being the probability simplex ∆ [ m ] [111, 112] we have a worst-case excess cost objective d C ( π, π π E ) =
max i ∈ [ m ] (cid:0)
ρ ϕ i ( π ) − ρ ϕ
i ( π E ) (cid:1) .
For clarity, we will replace c by w in the notation of the quantities defined in Section 2.1.
A Q -Convex-Analytic Viewpoint
Our methodology builds upon the convex-analytic approach to AL, first introduced by [112], with the key difference that we consider a different convex formulation that introduces Q -functions as slack variables.
This allows to design a practical scalable model-free algorithm with theoretical guarantees.
Let F ≜ { µ ∈ R |S||A| | ( B − γ P ) ⊺ µ = (1 − γ ) ν 0 , µ ≥ 0 } be the state-action polytope , where P is the vector form of P , i.e., P ( s,a ) ,s ′ ≜ P ( s ′ | s, a ) , and B is a binary matrix defined by B ( s,a ) ,
s ′ ≜if s = s ′ , and B ( s,a ) ,s ′ ≜ 0 otherwise.
The linear constraints that define the set F , also known as
Bellman flow constraints , precisely characterize the set of state-action occupancy measures.
Proposition 1 (94) .
We have that µ ∈ F if and only if there exists a unique stationary Markov policy π such that µ = µ π .
If µ ∈ F then the policy π µ ( a | x ) ≜ µ ( x,a ) (cid:80) a ′∈A µ ( x,a ′ ) has occupancy measure µ .
Using Proposition 1 and the dual form of the C -distance ¯ d C ( µ , µ π E ) = max w ∈W (cid:10) µ − µ π E , c w (cid:11) , it follows that (1) is equivalent to the primal convex program ζ
⋆ = min µ { ¯ d C ( µ , µ π E ) | µ ∈ F } .
In particular for W = ∆ [ m ] and by using an epigraphic transformation, we end up with an LP program [112], while for W =
B m 1 we get a quadratic objective with linear constraints [4].
A slight variation of the above reasoning is to introduce a mirror variable d and split the Bellman flow constraints in the definition of F .
We then get the primal convex program
[Equation]
(Primal) where the new polytope is given by M ≜ { ( µ , d )
| B ⊺ d = γ P ⊺ µ + (1 − γ ) ν 0 , µ = d , d ≥ 0 } .
This overparameterization trick has been first introduced by Mehta and Meyn [76] and has been recently revisited by [14, 84, 67, 83, 77, 71].
A salient feature of this equivalent formulation is that it introduces a Q -function as Lagrange multiplier to the equality constraint d = µ , and so lends itself to data-driven algorithms.
To motivate further this new formulation, in Appendix C, we shed light to its dual and provide an interpretation of the dual optimizers.
In particular, when W = B m 1 , we show that ( V ⋆ w true , Q ⋆ w true , w true ) is a dual optimizer.
For our theoretical analysis we focus on the linear MDP setting [55], i.e., we assume that the transition law is linear in the feature mapping.
We denote by ϕ ( s, a ) the ( s, a ) -th row of Φ .
Assumption 1 (Linear MDP) .
There exists a collection of m probability measures ω
= ( ω 1 , . .
. , ω m ) on S , such that P ( ·| s, a ) = ⟨ ω ( · ) , ϕ ( s, a ) ⟩ , for all ( s, a ) .
Moreover ϕ
( s, a ) ∈ ∆ [ m ] , for all ( s, a ) .
Assumption 1 essentialy says that the transition matrix P has rank at most m , and P = ΦM for some matrix M ∈ R m ×|S| .
It is worth noting that in the case of continuous MDPs, despite being linear, the transition law P ( ·| s, a ) can still have infinite degrees of freedom.
This is a substantial difference from the recent theoretical works on IL [70, 105] which consider either a linear quadratic regulator, or a transition law that can be completely specified by a finite-dimensional matrix such that the degrees of freedom are bounded.
Assumption 1 enables us to consider a relaxation of (Primal).
In particular, we aggregate the constraints µ = d by imposing Φ ⊺ µ = Φ ⊺ d instead, and introduce a variable λ =
Φ ⊺ µ .
It follows that λ lies in the m -dimensional simplex ∆ [ m ] .
Then, we get the following convex program
[Equation]
where M Φ ≜ { ( λ , d ) | B ⊺ d = γ M ⊺ λ + (1 − γ )
ν 0 , λ =
Φ ⊺ d , λ ∈ ∆ [ m ] , d ∈ ∆ S×A } .
As shown in [84, 14, 83], for linear MDPs, the set of occupancy measures F can be completely characterized by the set M Φ (c.f., Proposition 2).
While the number of constraints and variables in (Primal ′ ) is intractable for large scale MDPs, in the next paragraph, we show how this problem can be solved using a proximal point scheme.
Proximal Point Imitation Learning
By using a Lagrangian decomposition, we have that (Primal ′ ) is equivalent to the following bilinear saddle-point problem
[Equation]
where A ∈ R (2 m + |S| ) × ( m + |S||A| ) , and b ∈ R ( m + |S| + |S||A| ) are appropriately defined (see Appendix D), x ≜ [ λ ⊺ , d ⊺ ] ⊺ , y ≜ [ w ⊺ , V ⊺ , θ ⊺ ] ⊺ , X ≜ ∆ [ m ] × ∆ S×A , and Y ≜ W × R |S| × R m .
Since in practice we do not have access to the whole policy π E , but instead can observe a finite set of i.i.d.
sample trajectories D E ≜ { ( x ( l ) 0 , a ( l ) 0 , x ( l ) 1 , a ( l ) 1 , . . .
, x ( l ) H , a ( l ) H ) } n E l =1 ∼ π E , we define the vector (cid:98) b by replacing ρ Φ ( π E ) with its empirical counterpart
ρ Φ ( (cid:99) π E ) (by taking sample averages) in the definition of b .
We then consider the empirical objective f ( x ) ≜ max y ∈Y (cid:68) y , Ax + (cid:98) b (cid:69) and apply PPM on the decision variable x .
For the λ -variable we use the relative entropy D ( λ || λ ′ ) ≜ (cid:80) mi =1 λ ( i ) log λ ( i )
λ ′ ( i ) , while for the occupancy measure d we use the conditional relative entropy
H ( d || d ′ ) ≜ (cid:80) s,a d ( s, a ) log π d ( a | s ) π d ′ ( a | s ) .
With this choice we can rewrite the PPM update as ( λ k +1 , d k +1 ) =
[Equation]
where we used primal feasibility to replace λ k with Φ ⊺ d k as the center point of the relative entropy.
PPM is implicit, meaning that it requires the evaluation of the gradient at the next iterate x k +1 .
Such a requirement makes it not implementable in general.
However, in the following, we describe a procedure to apply proximal point to our specific f ( x ) .
The following Proposition summarizes the result.
Proposition 2.
For a parameter θ ∈ R m , we define the logistic state-action value function Q θ ∈
R |S||A| by Q θ ≜ Φ θ , and the k -step logistic state value function V k θ ∈ R |S| by
[Equation]
Moreover, we define the k -step reduced Bellman error function δ k w , θ ∈ R m by δ k w , θ ≜ w + γ MV k θ − θ .
Then, the PPM update ( λ ⋆k , d ⋆k ) in 2 is given by
[Equation]
where ( w ⋆k , θ ⋆k ) is the maximizer over W × R m of the k -step logistic policy evaluation objective
[Equation]
Moreover, it holds that G k
[Equation]
1 α
[Equation]
⋆k ||
d k − 1 ) .
If in addition Assumption 1 holds, then d ⋆k is a valid occupancy measure, i.e., d ⋆k ∈ F and so d ⋆k = µ π d ⋆k .
The proof of Proposition 2 is broken down into a sequence of lemmas and is presented in Appendix E.
It employs an analytical-oracle g : Y
→ X given by
[Equation]
and a max-oracle h : X → Y given by h ( x ) ≜ arg max y ∈Y ⟨ y , Ag ( y ; x ) ⟩ + 1 τ D Ω ( g ( y ; x ) || x ) , where we used D Ω to compact the two divergences.
By noting that the PPM update Equation (2) can be rewritten as x k +
=
g ( h ( x k ); x k ) , its analytical computation is reduced to the characterization of the two aforementioned oracles.
In particular, the updates (3)–(4) come from the analytical-oracle while (5) is the objective of the max-oracle .
The choice of conditional entropy as Bregman divergence for the λ variable living in the probability simplex is standard in the optimization literature and is known to mitigate the effect of dimension.
In particular, as noted in [85], the classic REPS algorithm [90] can be seen as mirror descent with relative entropy regularization.
On the other hand, the choice of conditional entropy as Bregman divergence for the d variable is less standard and has been popularized by Q-REPS [14].
Such particular divergence leads to an actor-critic algorithm that comes with several merits.
By Proposition 2, it is apparent that we get analytical softmin updates for the policy π d rather than the occupancy measure d .
Moreover, these softmin updates are expressed in terms of the logistic Q -function and do not involve the unknown transition matrix P .
Consequently, we avoid the problematic occupancy measure approximation and the restrictive coherence assumption on the choice of features needed in [13, 58], as well as the biased policy updates appearing in REPS [90, 89].
In addition, the newly introduced logistic policy evaluation objective G k ( w , θ ) has several desired properties.
It is concave and smooth in ( w , θ ) and has bounded gradients.
Therefore, it does not suffer from the pathologies of the squared
Bellman error [78] and does not require heuristic gradient clipping techniques.
Moreover, unlike [58] it allows a model-free implementation without the need for a generative model (see Section 4.1)
We stress the fact that the max-oracle of our proximal point scheme performs the cost update and policy evaluation phases jointly.
This is a rather novel feature of our algorithm that differs from the separate cost update and policy evaluation step used in recent theoretical imitation learning works [122, 105, 70].
Our joint optimization over cost and Q -functions avoids instability due to adversarial training and can also recover an explicit cost along with the Q -function without requiring knowledge or additional interaction with the environment (see Section 5).
It is worth noting that application of primal-dual mirror descent to (SPP) does not have this favorable property.
While in the standard MDP setting, proximal point and mirror descent coincide because of the linear objective, in imitation learning proximal point optimization makes a difference.
In Appendix K, we include a more detailed discussion and numerical comparison between PPM and mirror descent updates.
4.1 Practical Implementation
Exact optimization of the logistic policy evaluation objective is infeasible in practical scenarios, due to unknown dynamics and limited computation power.
In this section, we design a practical algorithm that uses only sample transitions by obtaining stochastic (albeit biased) gradient estimators.
Proposition 2 gives rise to Proximal Point Imitation Learning ( P 2 IL ), a model-free actor-critic IRL algorithm described in Algorithm 1.
The key feature of P 2 IL is that the policy evaluation step involves optimization of a single smooth and concave objective over both cost and state-action value function parameters.
In this way, we avoid instability or poor convergence in optimization due to nested policy evaluation and cost updates, as well as the undesirable properties of the widely used squared Bellman error.
In particular, the k th iteration of P 2 IL consists of the following two steps : i)
( Critic Step ) Computation of an approximate maximizer ( w k , θ k )
≈ arg max w , θ G k ( w , θ ) of the concave logistic policy evaluation objective, by using a biased stochastic gradient ascent subroutine; (ii) ( Actor Step ) Soft-min policy update π k ( a | s ) ∝ π k − 1 ( a | s ) e − αQ θ k
( s,a ) expressed in terms of the logistic Q -function.
The domain Θ in Algorithm 1 is the ℓ ∞ -ball with appropriately chosen radius D to be specified later (see Proposition 3).
Moreover, Π Θ ( x ) ≜ arg min y ∈ Θ ∥ x − y ∥ 2 (resp.
Π W ( w ) ) denotes the
Euclidean projection of x (resp.
w ) onto Θ (resp.
W ).
In order to estimate the gradients ∇ θ G k ( w , θ ) and ∇ w G k ( w , θ ) we invoke the Biased Stochastic
Gradient Estimator subroutine (BSGE) (Algorithm 2) given in Appendix H. By using the linear
MDP Assumption 1 and leveraging ridge regression and plug-in estimators, the proposed stochastic gradients can be computed via simple linear algebra with computational complexity poly ( m, n ( t )) , independent of the size of the state space.
Input: Feature matrix Φ , expert demonstrations D E , number of iterations K , step sizes η and α , number of SGD iterations T, SGD learning rates β
= { β t } T − 1 t =0 , number-of-samples function n : N → N
Initialize π 0 as uniform distribution over A
Compute the empirical FEV ρ Φ ( (cid:99) π E ) using expert demonstrations D E for k = 1 , . . .
K do //
Critic-step (policy evaluation)
Initialize θ k, 0 = 0 and w k, 0 =
Run π k − 1 and collect i.i.d.
samples B k = { ( s ( n ) k − 1 , a ( n ) k − 1 , s ′ ( n )
k − 1 ) } n ( T ) n =1 such that ( s ( n ) k − 1 , a ( n ) k − 1 ) ∼ µ π k − 1 and s ′ ( n )
k − 1 ∼ P ( ·| s ( n ) k − 1 , a ( n ) k − 1 ) for t = 0 , . . .
T − 1 do
Compute biased stochastic gradient estimators
[Equation]
θ
[Equation]
+
[Equation]
= Π
[Equation]
(cid:1) end for ( w
[Equation]
, θ
[Equation]
) = (
[Equation]
(cid:80)
[Equation]
1 w
[Equation]
(cid:80)
[Equation]
=1 θ
[Equation]
Actor-step (policy update)
Policy update: π k ( a | s ) ∝ π k − 1 ( a | s ) e − αQ θ k
( s,a ) end for
Output: Mixed policy (cid:98) π K of { π k } k ∈ [ K ]
4.2 Theoretical Analysis
The first step in our theoretical analysis is to study the propagation of optimization errors made by the algorithm on the true policy evaluation objective.
In particular at each iteration step k , the ideal policy evaluation update ( w ⋆k , θ ⋆k ) and the ideal policy update π ⋆k are given by ( w ⋆k , θ ⋆k ) =
arg max w , θ G k ( w , θ ) , and π ⋆k ( a | s ) =
π k − 1 ( a | s ) e − α
( Q θ ⋆k ( s,a ) − V k θ ⋆k ( s )) .
On the other hand, consider the realised policy evaluation update ( w k , θ k ) such that G k ( w ⋆k , θ ⋆k ) − G k ( w k , θ k )
=
ϵ k , the corresponding policy π k given by π k = π k − 1 ( a | s ) e − α
( Q θ k ( s,a ) − V k θ k ( s )) , and let d k ≜ µ π k .
We denote by (cid:98) π K the extracted mixed policy of { π k } Kk =1 .
We are interested in upper-bounding the suboptimality gap
d C ( (cid:98) π K , π E ) of Algorithm 1 as a function of ε k .
To this end, we need the following assumption.
Assumption 2.
It holds that λ min ( E ( s,a ) ∼ d k ϕ ( s, a ) ϕ
( s, a ) T ) ≥ β , for all k ∈
[ K ] .
Assumption 2 states that every occupancy measure d k induces a positive definite feature covariance matrix, and so every policy π k explores uniformly well in the feature space.
This assumption is common in the RL theory literature [2, 46, 37, 66, 3, 7].
It is also related to the condition of persistent excitation from the control literature [81].
The following proposition ensures that max
[Equation]
( w , θ ) = max w , θ ∈W× Θ G k ( w , θ ) .
Therefore, this constraint does not change the problem optimality, but will considerably accelerate the convergence of the algorithm by considering smaller domains.
Proposition 3.
There exists a maximizer θ ⋆k such that ∥ θ ⋆k ∥ ∞ ≤ 1+ | log β | 1 − γ ≜ D .
We can now state our error propagation theorem.
Theorem 1.
Let (cid:98) π K be the output of running Algorithm 1 for K iterations, with n E ≥ 2log( 2 m δ ) εexpert trajectories of length H ≥ 11 − γ log( 1 ε ) .
Let C ≜ 1 βη (cid:0) (cid:113) 2 α 1 − γ + √ 8 η (cid:1) + (cid:113) 18 α 1 − γ .
Then, with probability at least 1 − δ , it holds that d C ( (cid:98) π K , π E ) ≤ 1 K (cid:16) D ( λ ∗ ||
Φ T d 0 ) η
+ H
[Equation]
d
[Equation]
α
[Equation]
k
[Equation]
k
[Equation]
+ ε.
By Theorem 1, whenever the policy evaluation errors ε k , as well as the estimation error ε can be kept small, Algorithm 1 ouputs a policy (cid:98) π K with small suboptimality gap ρ c true ( (cid:98) π K ) − ρ c true ( π E ) .
[Figure]
e
[Figure]
Acrobot
[Figure]
GAIL
AIRL
[Figure]
Notably, there is no direct dependence on the size of the state space or the dimension of the feature space.
In the ideal case, where ε k = 0 for all k , the convergence rate is O (1 /K ) .
[-d e x      (a) − c true 0 1 
-]
The provided error propagation analysis still holds with general function approximation, i.e., in the context of deep RL.
Indeed, by choosing Φ =
I , Assumption 1 is trivially satisfied and the θ variable in the objective
G k is replaced by a Q -function.
In practice, the estimation error ε can be made arbitrary small, by increasing the number of expert demonstrations n E .
[-Suppose the trajectory length is H ≥ 
-]
Moreover, the next theorem ensures that under Assumptions 1 and 2 the biased stochastic gradient ascent (BSGA) subroutine has sublinear convergence rate.
Theorem 2.
Let ( w k , θ k ) be the output of the BSGA subroutine in Algorithm 1 for T iterations, with n ( t ) ≥ max (cid:16) O (cid:16) γ
[Equation]
mDt ( η + α ) 2 [-
-]β log Tmδ (cid:17) , O (cid:16) mt
[Equation]
(cid:17)(cid:17) sample transitions, and learning rates
β t
[-ni =
-]=
O ( 1 √ t ) .
[-The optimal choice for L is L = √ 
-]
Then, ϵ k = G k ( w ⋆k , θ ⋆k ) − G k ( w k , θ k ) ≤
O ( max { η, 1 } mD β √ T ) , with probability 1 − δ .
Corollary 1 (Resource guarantees) .
Choose η = α = 1 and let K
= Ω (cid:0) ϵ − 1 (cid:1) , T = Ω [-(cid:16) m 4 D 4 β 4 C 
-](cid:0) ϵ − 4 (cid:1) .
Then for Ω ( KT ) =
Ω (cid:0) ϵ − 5 (cid:1) sample transitions, Ω (cid:0) ε − 2 (cid:1) expert trajectories and approximately solving Ω (cid:0) ϵ − 1 (cid:1) concave maximization problems, we can ensure d C ( (cid:98) π, π E ) ≤
O ( ϵ + ε ) , with high probability.
Offline Setting.
Finally, we notice that using Φ ⊺ µ π E as the reference distribution for the relative entropy we can obtain an offline algorithm that does not require environment interactions.
By reinterpreting smoothing [82] as one step of proximal point, and using similar arguments as in the proof of Theorem 1, we can provide similar theoretical guarantees for the offline setting.
The formal statement of the theoretical result as well as the optimization of the empirical policy evaluation objective are presented in Appendix J (see Theorems 4 and 6).
Experiments
In this section, we demonstrate that our approach achieves convincing empirical performance in both online and offline IL settings on several environments.
The precise setting is detailed in Appendix L.
Online Setting.
We first present results in various tabular environments where we can implement our algorithm without any practical relaxation outperforming GAIL [51], AIRL [38] and IQ-Learn [40].
Results are given in Figure 1.
Good performance but inferior to IQ-Learn is observed also for continuous states environments (CartPole and Acrobot) where we used neural networks function approximation.
Offline Setting.
Figures 2a to 2c shows that our method is competitive with the state-of-the-art offline IL methods IQLearn [40] and AVRIL [25] that recently showed performances superior to other methods like [54][64].
We also tried our algorithm in the complex image-based Pong task from the Atari suite.
Figure 2d shows that the algorithm reaches the expert level after observing 2 e 5 expert samples.
We did not find AVRIL competitive in this setting, and skip it for brevity.
In these settings, we verified that the algorithmic performance is convincing even for costs parameterized by neural networks.
Continuous control experiments.
We attain the expert performance also in 2 MuJoCo environments:
Ant , HalfCheetah , Hopper , and Walker (see Figures 2e to 2h).
The additional difficulty in implementing the algorithm in continuous control experiments is that the analytical form of the policy improvement step is no longer computationally tractable because this would require to compute an integral over the continuous action space.
Therefore, we approximated this update using the Soft
Actor Critic (SAC) [44] algorithm.
SAC requires environment samples making the algorithm online.
The good empirical result opens the question of analyzing policy improvement errors as in [41].
[Figure]
Figure 2d shows the total return vs the number of
Recovered Costs.
A unique algorithmic feature of the proposed methodology is that we can explicitly recover a cost along with the Q -function without requiring adversarial training.
In Figure 3, we visualize our recovered costs in a simple 5x5 Gridworld .
Most importantly, we verify that the recovered costs induce nearly optimal policies w.r.t.
the unknown true cost function.
Compared to I Q -Learn [40], we do not require knowledge or further interaction with the environment.
Therefore, the recovered cost functions show promising transfer capability to new dynamics. [Figure] index S t a t e I n
{+d e x      (a) − c true 0 1 
 Action index S t+} a t e I n d e x      (b) − c K 0 1 2 3 4    4600470048004900(c)
- V ⋆ c true 0 1 2 3 4    232236240244(d) - V ⋆ c K Cost Transfer Setting. We experimented with a transfer cost setting on a Gridworld (Figure 4). We consider two different Gridworld MDP environments, say M and (cid:102) M , with opposite action effects.
This means that action Down in (cid:102)
M corresponds to action Left in M and vice versa.
Similarly, the effects of Up and Right are swapped between (cid:102) M and M .
We denote by V π (cid:102)
M, c true (resp.
V ⋆ (cid:102) M, c true ) the value function of policy π (resp.
optimal value function) in the MDP environment (cid:102)
M with cost function c true .
Moreover, we denote by π ⋆M, c
the optimal policy in the MDP environment M under cost function c .
Figure (a) gives the corresponding optimal value function.
Figure (b) presents the value function of the expert policy π E = π ⋆M, c true used as target by P 2 IL . Figure (d) shows the value function of the learned imitating policy π K from P 2 IL . Finally, Figure (b) depicts the value function of the optimal policy π ⋆ (cid:102) M, c K for the environment (cid:102)
M endowed with the recovered cost function c K by
P 2 IL (with access to samples from M ).
We conclude that the policy π ⋆ (cid:102) M, c K is optimal in (cid:102) M with cost c true . By contrast, the expert policy π E = π ⋆M, c true used as target by P 2 IL performs poorly and as a consequence also the imitating policy π K does so.
All in all, we notice that the recovered cost induces an optimal policy for the new dynamics while the imitating policy fails. Albeit, cost transfer is successful in this experiment we do not expect this fact to be true in general because we do not tackle the issue of cost shaping [87].
[Figure]
Discussion and Outlook
In this work, we studied a Proximal Point Imitation Learning ( P 2 IL ) algorithm with both theoretical guarantees and convincing empirical performance.
Our methodology is rooted in classical optimization tools and the LP approach to MDPs. The most significant merits of P 2 IL are the following: (i) It optimizes a convex and smooth logistic Bellman evaluation objective over both cost and Q-functions. In particular, it avoids instability due to adversarial training and can also recover an explicit cost along with Q function; (ii)
In the context of linear MDPs, it comes with efficient resource guarantees and error bounds for the suboptimality of the learned policy (Theorem 2 and Corollary 1). In particular, given poly(1 /ε, log(1 /δ ) , m ) many samples , it recovers an ε -optimal policy, with probability 1 − δ . Notably, the bound is independent of the size of the state-action space; (iii) Beyond the linear MDP setting, it can be implemented in a model-free manner, for both online and offline setups, with general function approximation without losing its theoretical specifications.
This is justified by providing an error propagation analysis (Theorems 1 and 4), guaranteeing that small optimization errors lead to high-quality output policy; (iv) It enjoys not only strong theoretical guarantees but also favorable empirical performance. At the same time, our newly introduced methods bring challenges and open questions.
One interesting question is whether one can accelerate the PPM updates and improve the convergence rate.
Another direction for future work is to provide rigorous arguments for the near-optimality of the recovered cost function.
On the practical side, we plan to conduct experiments in more challenging environments than MuJoCo and Atari. We hope our new techniques will be useful to future algorithm designers and lay the foundations for overcoming current limitations and challenges.
In Appendix B, we point out in detail a few interesting future directions.
The authors would like to thank the anonymous reviewer for their suggestions to improve the presentation and for motivating us to inspect the recovered cost function. This work has received funding from the Enterprise for Society Center (E4S), the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme grant agreement OCAL, No. 787845, the European Research Council (ERC) under the European Union’s Horizonresearch and innovation programme (grant agreement n° 725594 - time-data), the Swiss National
Science Foundation (SNSF) under grant number 200021_205011. Gergely Neu was supported by the European Research Council (ERC) under the European Union’s Horizon 2020 research and innovation programme (Grant agreement No. 950180).
Luca Viano acknowledges travel support from ELISE (GA no 951847). Checklist For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes] We reflect the contribution described in the introduction with the theoretical results in Section 4, Section 4.
(b) Did you describe the limitations of your work?
[Yes] We state the assumptions needed for the analysis in Sections 3 and 4. (c) Did you discuss any potential negative societal impacts of your work?
[N/A]
The work is mainly theoretical, we do not foresee potential negative impacts on the society. (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] We acknowledge habing read the review guidelines. If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] We state the assumptions needed for the analysis in Section 3 Section 4. (b) Did you include complete proofs of all theoretical results? [Yes] The main results are stated in Section 4.2 and the proofs are included as supplementary material. If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] The code is included in the supplementary material. (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Training details are provided in the Appendix. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] We averaged multiple seeds whenever possible computationally.
We ran a single seed for the computational expensive Pong environment. (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] We specified the resource in the Supplementary material. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] We cite [40] for using their codebase and their expert data.
(b) Did you mention the license of the assets? [Yes] We mention the license in the Supplementary. (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] We attach the code to the supplementary.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [Yes] We discuss this in the supplemnetary material (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A]
We do not use personal data. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
We did not involve human partecipants. (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A]
We did not involve human partecipants. (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
We did not involve human partecipants.