Hamiltonian Q-Learning: Leveraging Importance-sampling for Data Efficient RL

Model-free reinforcement learning ( RL ) , in particular Q-learning is widely used to learn optimal policies for a variety of planning and control problems . However , when the underlying state-transition dynamics are stochastic and high-dimensional , Q-learning requires a large amount of data and incurs a prohibitively high computational cost . In this paper , we introduce Hamiltonian Q-Learning , a data efficient modification of the Q-learning approach , which adopts an importance-sampling based technique for computing the Q function . To exploit stochastic structure of the state-transition dynamics , we employ Hamiltonian Monte Carlo to update Q function estimates by approximating the expected future rewards using Q values associated with a subset of next states . Further , to exploit the latent low-rank structure of the dynamic system , Hamiltonian Q-Learning uses a matrix completion algorithm to reconstruct the updated Q function from Q value updates over a much smaller subset of state-action pairs . By providing an efficient way to apply Qlearning in stochastic , high-dimensional problems , the proposed approach broadens the scope of RL algorithms for real-world applications , including classical control tasks and environmental monitoring . 1 INTRODUCTION . In recent years , reinforcement learning ( Sutton & Barto , 2018 ) have achieved remarkable success with sequential decision making tasks especially in complex , uncertain environments . RL algorithms have been widely applied to a variety of real world problems , such as resource allocation ( Mao et al. , 2016 ) , chemical process optimization ( Zhou et al. , 2017 ) , automatic control ( Duan et al. , 2016 ) , and robotics ( Kober et al. , 2013 ) . Existing RL techniques often offer satisfactory performance only when it is allowed to explore the environment long enough and generating a large amount of data in the process ( Mnih et al. , 2015 ; Kamthe & Deisenroth , 2018 ; Yang et al. , 2020a ) . This can be prohibitively expensive and thereby limits the use of RL for complex decision support problems . Q-Learning ( Watkins , 1989 ; Watkins & Dayan , 1992 ) is a model-free RL framework that captures the salient features of sequential decision making , where an agent , after observing current state of the environment , chooses an action and receives a reward . The action chosen by the agent is based on a policy defined by the state-action value function , also called the Q function . Performance of such policies strongly depends on the accessibility of a sufficiently large data set covering the space spanned by the state-action pairs . In particular , for high-dimensional problems , existing model-free RL methods using random sampling techniques leads to poor performance and high computational cost . To overcome this challenge , in this paper we propose an intelligent sampling technique that exploits the inherent structures of the underlying space related to the dynamics of the system . It has been observed that formulating planning and control tasks in a variety of dynamical systems such as video games ( Atari games ) , classical control problems ( simple pendulum , cart pole and double integrator ) and adaptive sampling ( ocean sampling , environmental monitoring ) as Q-Learning problems leads to low-rank structures in the Q matrix ( Ong , 2015 ; Yang et al. , 2020b ; Shah et al. , 2020 ) . Since these systems naturally consist of a large number of states , efficient exploitation of low rank structure of the Q matrix can potentially lead to significant reduction in computational complexity and improved performance . However , when the state space is high-dimensional and further , the state transition is probabilistic , high computational complexity associated with calculating the expected Q values of next states renders existing Q-Learning methods impractical . A potential solution for this problem lies in approximating the expectation of Q values of next states with the sample mean of Q values over a subset of next states . A natural way to select a subset of next states is by drawing IID samples from the transition probability distribution . However , this straight forward approach becomes challenging when the state transition probability distribution is highdimensional and is known only up to a constant . We address this problem by using Hamilton Monte Carlo ( HMC ) to sample next states ; HMC draws samples by integrating a Hamiltonian dynamics governed by the transition probability ( Neal et al. , 2011 ) . We improve the data efficiency further by using matrix completion methods to exploit the low rank structure of a Q matrix . RELATED WORK Data efficient Reinforcement Learning : The last decade has witnessed a growing interest in improving data efficiency in RL methods by exploiting emergent global structures from underlying system dynamics . Deisenroth & Rasmussen ( 2011 ) ; Pan & Theodorou ( 2014 ) ; Kamthe & Deisenroth ( 2018 ) ; Buckman et al . ( 2018 ) have proposed model-based RL methods that improve data efficiency by explicitly incorporating prior knowledge about state transition dynamics of the underlying system . Dearden et al . ( 1998 ) ; Koppel et al . ( 2018 ) ; Jeong et al . ( 2017 ) propose Baysean methods to approximate the Q function . Ong ( 2015 ) and Yang et al . ( 2020b ) consider a model-free RL approach that exploit structures of state-action value function . The work by Ong ( 2015 ) decomposes the Q matrix into a low-rank and sparse matrix model and uses matrix completion methods ( Candes & Plan , 2010 ; Wen et al. , 2012 ; Chen & Chi , 2018 ) to improve data efficiency . A more recent work by Yang et al . ( 2020b ) has shown that incorporation of low rank matrix completion methods for recovering Q matrix from a small subset of Q values can improve learning of optimal policies . At each time step the agent chooses a subset of state-action pairs and update their Q value according to the Bellman optimally equation that considers a discounted average between reward and expectation of the Q values of next states . Shah et al . ( 2020 ) extends this work by proposing a novel matrix estimation method and providing theoretical guarantees for the convergence to a -optimal Q function . On the other hand , entropy regularization ( Ahmed et al. , 2019 ; Yang et al. , 2019 ; Smirnova & Dohmatob , 2020 ) , by penalizing excessive randomness in the conditional distribution of actions for a given state , provides an alternative means to implicitly exploit the underlying low-dimensional structure of the value function . Lee et al . ( 2019 ) proposes an approach that samples a whole episode and then updates values in a recursive , backward manner . CONTRIBUTION The main contribution of this work is three-fold . First , we introduce a modifiedQ-learning framework , called Hamiltonian Q-learning , which uses HMC sampling for efficient computation of Q values . This innovation , by proposing to sample Q values from the region with the dominant contribution to the expectation of discounted reward , provides a data-efficient approach for using Q-learning in real-world problems with high-dimensional state space and probabilistic state transition . Furthermore , integration of this sampling approach with matrix-completion enables us to update Q values for only a small subset of state-action pairs and thereafter reconstruct the complete Q matrix . Second , we provide theoretical guarantees that the error between the optimal Q function and the Q function obtained by updating Q values using HMC sampling can be made arbitrarily small . This result also holds when only a handful ofQ values are updated using HMC and the rest are estimated using matrix completion . Further , we provide theoretical guarantee that the sampling complexity of our algorithm matches the mini-max sampling complexity proposed by Tsybakov ( 2008 ) . Finally , we demonstrate the effectiveness of Hamiltonian Q-learning by applying it to a cart-pole stabilization problem and an adaptive ocean sampling problem . Our results also indicate that our proposed approach becomes more effective with increase in state space dimension . 2 PRELIMINARY CONCEPTS . In this section , we provide a brief background on Q-Learning , HMC sampling and matrix completion , as well as introduce the mathematical notations . In this paper , |Z| denotes the cardinality of a set Z . Moreover , R represent the real line and AT denotes the transpose of matrix A . 2.1 Q-LEARNING Markov Decision Process ( MDP ) is a mathematical formulation that captures salient features of sequential decision making ( Bertsekas , 1995 ) . In particular , a finite MDP is defined by the tuple ( S , A , P , r , γ ) , where S is the finite set of system states , A is the finite set of actions , P : S×A×S → [ 0 , 1 ] is the transition probability kernel , r : S ×A → R is a bounded reward function , and γ ∈ [ 0 , 1 ) is a discounting factor . Without loss of generality , states s ∈ S and actions a ∈ A can be assumed to be Ds-dimensional and Da-dimensional real vectors , respectively . Moreover , by letting si denote the ith element of a state vector , we define the range of state space in terms of the following intervals [ d−i , d + i ] such that s i ∈ [ d−i , d+i ] ∀i ∈ { 1 , . . . , Ds } . At each time t ∈ { 1 , . . . , T } over the decision making horizon , an agent observes the state of the environment st ∈ S and takes an action at according to some policy π which maximizes the discounted cumulative reward . Once this action has been executed , the agent receives a reward r ( st , at ) from the environment and the state of the environment changes to st+1 according to the transition probability kernel P ( ·|st , at ) . The Q function , which represents the expected discounted reward for taking a specific action at the current time and following the policy thereafter , is defined as a mapping from the space of state-action pairs to the real line , i.e . Q : S × A → R. Then , by letting Qt represent the Q matrix at time t , i.e . the tabulation of Q function over all possible state-action pairs associated with the finite MDP , we can express the Q value iteration over time steps as Qt+1 ( st , at ) = ∑ s∈S P ( s|st , at ) ( r ( st , at ) + γmax a Qt ( s , a ) ) . ( 1 ) Under this update rule , the Q function converges to its unique optimal value Q∗ ( Melo , 2001 ) . But computing the sum ( 1 ) over all possible next states is computationally expensive in certain problems ; in these cases taking the summation over a subset of the next states provides an efficient alternative for updating the Q values . 2.2 HAMILTONIAN MONTE CARLO . Hamiltonian Monte Carlo is a sampling approach for drawing samples from probability distributions known up to a constant . It offers faster convergence than Markov Chain Monte Carlo ( MCMC ) sampling ( Neal et al. , 2011 ; Betancourt ; Betancourt et al. , 2017 ; Neklyudov et al. , 2020 ) . To draw samples from a smooth target distribution P ( s ) , which is defined on the Euclidean space and assumed to be known up to a constant , HMC extends the target distribution to a joint distribution over the target variable s ( viewed as position within the HMC context ) and an auxiliary variable v ( viewed as momentum within the HMC context ) . We define the Hamiltonian of the system as H ( s , v ) = − logP ( s , v ) = − logP ( s ) − logP ( v|s ) = U ( s ) +K ( v , s ) , where U ( s ) , − logP ( s ) and K ( v , s ) , − logP ( v|s ) = 12vTM−1v represent the potential and kinetic energy , respectively , and M is a suitable choice of the mass matrix . HMC sampling method consists of the following three steps − ( i ) a new momentum variable v is drawn from a fixed probability distribution , typically a multivariate Gaussian ; ( ii ) then a new proposal ( s′ , v′ ) is obtained by generating a trajectory that starts from ( s , v ) and obeys Hamiltonian dynamics , i.e . ṡ = ∂H∂v , v̇ = −∂H∂s ; and ( iii ) finally this new proposal is accepted with probability min { 1 , exp ( H ( s , v ) −H ( s′ , −v′ ) ) } following the Metropolis–Hastings acceptance/rejection rule . 2.3 LOW-RANK STRUCTURE IN Q-LEARNING AND MATRIX COMPLETION Prior work ( Johns & Mahadevan , 2007 ; Geist & Pietquin , 2013 ; Ong , 2015 ; Shah et al. , 2020 ) on value function approximation based approaches for RL has implicitly assumed that the state-action value functions are low-dimensional and used various basis functions to represent them , e.g . CMAC , radial basis function , etc . This can be attributed to the fact that the underlying state transition and reward function are often endowed with some structure . More recently , Yang et al . ( 2020b ) provide empirical guarantees that the Q-matrices for benchmark Atari games and classical control tasks exhibit low-rank structure . Therefore , using matrix completion techniques ( Xu et al. , 2013 ; Chen & Chi , 2018 ) to recover Q ∈ R|S|×|A| from a small number of observed Q values constitutes a viable approach towards improving data efficiency . As low-rank matrix structures can be recovered by constraining the nuclear norm , the Q matrix can be reconstructed from its observed values ( Q̂ ) by solving Q = arg min Q̃∈R|S|×|A| ‖Q̃‖∗ subject to JΩ ( Q̃ ) = JΩ ( Q̂ ) , ( 2 ) where ‖ · ‖∗ denotes the nuclear norm ( i.e. , the sum of its singular values ) , Ω is the observed set of elements , and JΩ is the observation operator , i.e . JΩ ( x ) = x if x ∈ Ω and zero otherwise . 3 HAMILTONIAN Q-LEARNING A large class of real world sequential decision making problems - for example , board/video games , control of robots ’ movement , and portfolio optimization - involves high-dimensional state spaces and often has large number of distinct states along each individual dimension . As using a Q-Learning based approach to train RL-agents for these problems typically requires tens to hundreds of millions of samples ( Mnih et al. , 2015 ; Silver et al. , 2017 ) , there is a strong need for data efficient algorithms forQ-Learning . In addition , state transition in such systems is often probabilistic in nature ; even when the underlying dynamics of the system is inherently deterministic ; presence of external disturbances and parameter variations/uncertainties lead to probabilistic state transitions . Learning an optimal Q∗ function through value iteration methods requires updating Q values of state-action pairs using a sum of the reward and a discounted expectation of Q values associated with next states . In this work , we assume the reward to be a deterministic function of state-action pairs . However , when the reward is stochastic , these results can be extended by replacing the reward with its expectation . Subsequently , we can express ( 1 ) as Qt+1 ( st , at ) = r ( st , at ) + γE ( max a Qt ( s , a ) ) , ( 3 ) where E denotes the expectation over the discrete probability measure P. When the underlying state space is high-dimensional and has large number of states , obtaining a more accurate estimate of the expectation is computationally very expensive . The complexity increases quadratically with the number of states and linearly with number of actions , rendering the existing algorithms impractical . In this work , we propose a solution to this issue by introducing an importance-sampling based method to approximate the aforementioned expectation from a sample mean of Q values over a subset of next states . A natural way to sample a subset from the set of all possible next states is to draw identically and independently distributed ( IID ) samples from the transition probability distribution P ( ·|st , at ) . However , when the transition probability distribution is high-dimensional and known only up to a constant , drawing IID samples incurs a very high computation cost .

Summary: This work focuses on dynamic programming in the tabular setting. It proposes to use Hamiltonian Monte-Carlo (HMC) to sample the next states (instead of IID samples) and matrix completion to learn a low-rank Q matrix. It shows theoretical convergence. Experiments on discretized problems (CartPole and an ocean sampling problem) show that HMC and low-rank learning can behave more benignly compared to IID samples.