Lagrangian Method for Episodic Learning

1 INTRODUCTION . Episodic learning Terry ( 2017 ) is a general learning paradigm in which the agent learns based on data collected from a sequence of episodes of environmental interactions . Each episode consists of a finite number of decision steps , and the goal is to maximize an episode-wise performance . With one-step episodes , episodic learning subsumes supervised learning and bandit problems as special cases ; for sequential decision making problems , an episode may contain multiple ( and possibly varying number of ) steps , in which case episode learning encompasses most reinforcement learning and imitation learning problems encountered in current AI practice . A pervasive idea of episodic learning , often called the value-based approach Sutton and Barto ( 2018 ) , seeks to find a value function such that greedily selecting the actions with the best values ( as prescribed by this function ) gives good performance . For example , in most classification tasks , most deep neural networks are learned ( and used ) to represent such a value function . Finding good value functions can be highly challenging , especially for tasks with multi-step episodes . In this case , the classic Bellman optimality equation Bertsekas ( 2019 ) gives a sufficient condition of optimal value function ( i.e . value functions with which greedy decisions have optimal performance ) . In this paper , we study a Lagrangian method which learns optimal value functions by characterizing them as saddle points of the Lagrangian function of a variational ( re- ) formulation of the Bellman equation . The method reveals an elegant thoery on policy-value duality in machine learning . A restricted form of this high-level idea has been known as the Linear Programming ( LP ) approach in the literature of dynamic programming Puterman ( 1994 ) ; Si et al . ( 2004 ) ; Chen and Wang ( 2016 ) ; Wang ( 2017 ) , which has also been introduced to reinforcement learning in recent years Dai et al . ( 2018 ) ; Lee and He ( 2018 ) ; Nachum et al . ( 2019 ) ; Yang et al . ( 2020 ) ; Nachum and Dai ( 2020 ) . This paper provides a new perspective to the foundation of this elegant idea through the following contributions : Firstly , previous works are mostly limited to the discounted-reward setting , which bypass some technical difficulties Nachum and Dai ( 2020 ) , but deviate from the practical settings of many ( if not most ) real-world tasks . This paper fills this well-known gap between theory and practice by establishing a value-based theory ( and algorithm ) that targets at directly learning optimal value functions and policies for the undiscounted episode-wise reward ( see ( 1 ) in Section 2 ) . Secondly , previous works focused on linear formulations , which enjoys the strong duality property and generic LP techniques , but are consequently limited to treatments for the state-value functions ( i.e . V-functions ) in policy optimization setting . However , the optimal state-value function can not be directly used to derive optimal policy in model-free learning . In contrast , we proved a minimax theorem for the nonlinear Lagrangian functions associated with the action-value functions ( i.e . Qfunctions ) . The theorem opens a door to nonlinear variational treatment to the Bellman equation , which now enjoys the strong duality property too ( see Section 3 ) . In particular , we developed a simple imitation learning algorithm based on the Lagrangian duality thus established , and applied the algorithm to Machine Translation ( MT ) as case study . Transformer models trained by our algorithm achieves beam-search-level performance Vaswani et al . ( 2017 ) with only greedy decoding , and leads to 1.4-BLEU improvement when also equipping with beam search ( see Section 4 ) . Last but not the least , previous works have exclusively focused on the primal-form optimization formulation of the Bellman equation , thus are limited to solving the minimax-type saddle points of the corresponding Lagrangian function . We however showed that the minimax saddle points are not necessarily optimal value functions in learning settings . Instead , another class of the Lagrangian saddle points – the maximin-type saddle points which are derived from the dual-form optimization formulation – turn out to have rigorous guarantee on optimality . This observation points to new directions for this decades-old topic ( see Section 5 ) . 2 EPISODIC LEARNING PROCESS AND BELLMAN OPTIMAL VALUE . We start with introducing the mathematical formulation of episode learning problems , as well as the formal definitions of related concepts . An infinite-horizon MDP is a tuple ( S , A , R , P , ρ0 ) , where S is the state space , A the action space , R ( s ) ∈ [ rmin , rmax ] is a bounded reward ( possibly negative ) associated to each state s ∈ S , 1 P ( s′|s , a ) specifies action-conditioned transition probabilities between states , and ρ0 is the initial distribution with S0 ∼ ρ0 . A MDP is finite if both S and A are finite sets . A policy function π : S × A → [ 0 , 1 ] specifies the action selection probabilities under each state , which induces a markov chain Pπ [ St+1 = s′|St = s ] . = ∑ a∈A P ( s ′|s , a ) · π ( a|s ) . Let Π denote the policy space , i.e. , the set of all policies . Without loss of generality , we assume every state s in the state space S is reachable from the initial state under at least one policy , 2 where reachable means ∃π ∈ Π , ∑∞ t=0 Pπ [ St = s ] > 0 . An Episodic Learning Process ( ELP ) White ( 2017 ) ; Bojun ( 2020 ) is an infinite-horizon MDP that repeatedly encounters into , and is reset by , a group of terminal states S⊥ ( despite its name , a terminal state does not terminate the ELP ) . Formally , an infinite-horizon MDP is an ELP if there is a non-empty subset S⊥ ⊆ S such that ( 1 ) all terminal states have homogeneous and action-agnostic outbound probabilities : P ( s′|s1 , a1 ) = P ( s′|s2 , a2 ) , ∀s1 , s2 ∈ S⊥ , ∀a1 , a2 ∈ A , ∀s′ ∈ S , ( 2 ) the initial state is a terminal state : ρ0 ( s ) = 0 , ∀s 6∈ S⊥ , and ( 3 ) the average episode length is finite under any policy : ∀π ∈ Π , Eπ [ T ] < ∞ , where T . = inf { t ≥ 1 : St ∈ S⊥ } is called the termination time . Bojun ( 2020 ) proved that in any ELP , every policy π ∈ Π has a unique stationary distribution ρπ ( s , a ) such that Pπ [ St+1 = s , At+1 = a ] = ρπ ( s , a ) if Pπ [ St = s , At = a ] = ρπ ( s , a ) . Following White ( 2017 ) and Bojun ( 2020 ) , we use the ELP formalism to study finite-time decision tasks , which arguably account for most AI tasks encountered in current practice . The ELP model formulates the learning process of such task , which consists of an infinite number of concatenated episodes , each starting and ending at an terminal state . In general , an episode can be arbitrarily long , but will terminate in finite steps with probability 1 . The learning agent may choose to use a different behavior policy βt for each step t of the episodic learning process ( based on data collected from 1Our state-based reward formulation follows Schulman et al . ( 2015 ) and Bojun ( 2020 ) , and is equivalent to action-based reward formulations in terms of expressiveness . Refer to Appendix A of Bojun ( 2020 ) for details . 2Unreachable states are irrelevant to the real world whatsoever , so are excluded from our treatment . previous steps ) , and the goal is to have βt converge ( as t goes to infinity ) towards an optimal policy that maximizes the expected episode-wise cumulative reward : J ( π ) . = E ζ∼π [ ∑T ( ζ ) t=1 R ( St ) ] ( 1 ) where ζ = ( S0 , S1 , S2 , . . . ) is an infinite trajectory of the markov chain induced by policy π . Note that the termination time T is a random variable whose value may vary for different trajectory ζ . See Section D.3 for the ELP formulation of machine translation as an example . A value function Q : S × A → R assigns a real number to each state-action pair ( s , a ) as the “ perceived benefit ” of doing action a under state s. We call the set of all possible value functions , the value space , denoted by Q . A value function Q ∈ Q is called an optimal value if ( any of the ) Q-greedy policy with πQ ( a|s ) > 0 ⇒ Q ( s , a ) = maxāQ ( s , ā ) is an optimal policy . Bellman optimal value is a special optimal value function that is characterized by the Bellman optimality operator . In its general form Degris et al . ( 2012 ) ; Sutton et al . ( 2011 ; 2014 ; 2016 ) , a generalized Bellman optimality operator Bγ : Q → Q transforms a value function Q ∈ Q into another value function BγQ ∈ Q such that for any ( s , a ) ∈ S ×A , Bγ Q ( s , a ) .= ∑ s′∈S P ( s′|s , a ) · ( R ( s′ ) + γ ( s′ ) ·max a′∈A Q ( s′ , a′ ) ) ( 2 ) where γ : S → [ 0 , 1 ] is a discounting function over the states . As a classic result , if the discounting function equals a constant γc < 1 on every state , the corresponding Bellman optimality operator Bγc has a unique fixed point in the value space of any MDP . However , the fixed point of Bγc is generally not an optimal value with respect to the undiscounted objective ( 1 ) . In the following , we present a new theorem that guarantees the uniqueness of Bellman fixed point for a more general class of discounting functions , among which a particular “ episodic discounting function ” gives optimal value w.r.t . objective ( 1 ) . Both the uniqueness and the optimality are fundamentally rooted from an inherent graph property of ELP-MDPs . See the proof in Appendix A for detailed elaboration . Theorem 1 . In any finite Episodic Learning Process ( S , A , P , R , ρ0 ) , let γ be any discounting function such that γ ( s ) < 1 for all terminal state s ∈ S⊥ , then ( 1 ) Bγ has a unique fixed point , i.e. , the equation Q = BγQ has a unique solution . ( 2 ) The fixed point of Bγ is the limiting point of repeatedly applying Bγ to any Q ∈ Q . ( 3 ) The fixed point of Bγ is an optimal value w.r.t . objective ( 1 ) if γ is the following episodic discounting function : γ epi ( s ) . = 1 [ s 6∈ S⊥ ] = { 1 for s 6∈ S⊥ 0 for s ∈ S⊥ . ( 3 ) Since we focus on optimizing objective ( 1 ) , in the rest of the paper : Bellman optimality operator , denoted by B .= Bγ epi , refers to the specific operator ( 2 ) that uses the particular episodic discounting function ( 3 ) ; accordingly , Bellman optimal value , denoted by Q∗ , refers to the unique fixed point of B ; and similarly , Bellman optimality equation refers to the fixed-point equation Q = BQ under episodic discounting , or more explicitly , refers to the following system of non-linear equations : Q ( s , a ) = ∑ s′∈S max a′∈A P ( s′|s , a ) · ( R ( s′ ) + γ epi ( s ′ ) ·Q ( s′ , a′ ) ) , ∀ ( s , a ) ∈ S ×A ( 4 ) It is worthwhile noting that although the Bellman optimal value function is unique , there can be many optimal value functions in an episodic learning problem . In particular , any value function that gives the same preference order with the Bellman optimality value is an optimal value function .

Summary: In this work, the authors consider nonlinear Q-form Lagrangian function and show corresponding strong duality property. The main contributions are 1) the new proof of showing the duality gap as zero from a minmax perspective; 2) the generality of the theory with applications to machine translation tasks. A imitation learning algorithm is proposed as well and it turns out that it works well compared with the existing one.