Policy Augmentation: An Exploration Strategy for Faster Convergence of Deep Reinforcement Learning Algorithms

Despite advancements in deep reinforcement learning algorithms, developing an effective exploration strategy is still an open problem. Most existing exploration strategies either are based on simple heuristics, or require the model of the environment, or train additional deep neural networks to generate imagination-augmented paths. In this paper, a revolutionary algorithm, called Policy Augmentation, is introduced. Policy Augmentation is based on a newly developed inductive matrix completion method. The proposed algorithm augments the values of unexplored state-action pairs, helping the agent take actions that will result in high-value returns while the agent is in the early episodes. Training deep reinforcement learning algorithms with high-value rollouts leads to the faster convergence of deep reinforcement learning algorithms. Our experiments show the superior performance of Policy Augmentation. The code can be found at: https://github.com/arashmahyari/PolicyAugmentation.


INTRODUCTION
Reinforcement learning (RL) has shown significant success in recent years [1,2,3,4,5]. RL studies how an agent can maximize its cumulative rewards in a previously unknown environment through experience. RL algorithms maximize their long-term cumulative rewards by exploring new states, and their short-term (immediate) cumulative rewards by exploiting their learned policy.
The main challenge with RL algorithms is balancing between exploration and exploitation. While exploitation is critical to maximize the agent's immediate reward, the exploration ensures that the agent's overall (long-term) reward is maximized. Agents must explore most states to learn which states are more rewarding. However, the random exploration of the environment leads to a long convergence time. Depending on the problem and the environment, RL algorithms may require millions of episodes to explore all states and learn the policy. The -greedy algorithm-a simple exploration strategy-takes either random actions (exploration) or actions according to the policy (exploitation). The simple strategies such as -greedy well behave in situations where the reward is well-shaped. In more complicated tasks with sparse rewards, these strategies are not efficient.
More recently, several exploration strategies, e.g., curiositydriven exploration [6], count-based exploration [7], etc., have been proposed to improve the exploration of RL algorithms. The Uniform sampling strategy draws a random sample uniformly from the stored samples to update the action-value function-Q functions [2]. Trust region policy optimization (TRPO) was proposed to generate a number of trajectories and randomly select a subset of states from these trajectories [8]. The count-based exploration strategy [7] encourages agents to explore the less-visited states. These algorithms fail to perform optimally for tasks with sparse rewards. Besides, they simulate the future returns of rollout sets according to the current approximated policy. The performance of these algorithms depends on the approximated action-value functions according to previously explored state-action pairs. While these strategies are different from each other, they share a common idea that agents must be encouraged to visit less-explored states. These strategies take actions randomly from less-visited states, which may or may not lead to a reward. The existing exploration strategies do not encourage exploring state-action pairs that will return higher values.
In this paper, we propose Policy Augmentation to predict the values of unexplored state-action pairs by augmenting the action-value function, Q, using Inductive Matrix Completion (IMC). Matrix completion recovers the underlying low-rank matrix from a given subset of its entities [9]. IMC extends this idea to the association inference prediction using the available side information about the columns and rows of the underlying low-dimensional matrix [10]. In this paper, IMC recovers the Q function from the limited observed entities-the limited number of explored state-action pairs-using the features of states and actions. Figure 1 shows an example of recovered Q function for the Mountain car environment. The agent, then, takes actions according to the recovered action-value functions-the augmented policy. Taking actions according to the recovered Q function results in generating high-value rollouts for training deep reinforcement learning (DRL) algorithms. The purpose of Policy Augmentation is to generate high-value rollouts and train the agent with them during early episodes, which reduces the training time significantly. After training with high-value rollouts, the agent explores and exploits according to the existing DRL approaches, e.g., Proximal Policy Optimization (PPO) [11]. Fig. 3.a shows the classical Mountain Car control problem [12]. The environment is described in details in Section 4. The optimal policy must accelerate to the right when the car moves in the direction of the positive velocity, and accelerate to the left when the car is moving in the direction of the negative velocity. The existing exploration methods may apply negative, positive, or zero force when the car moves in the direction of the positive velocity. Policy Augmentation explores the pairs that are predicted to have high values, e.g., positive force in the direction of the positive velocity. Our proposed strategy is different from [13], in that our method requires neither training nor the model of the environment, whereas the deep learning model used in [13] requires training.

REINFORCEMENT LEARNING
An infinite-horizon discounted Markov decision process (MDP) is defined by (S ; A ; P, r, ρ 0 , γ), where S is a set of finite states, A is a finite set of actions, P : S ×A ×S → R is the transition probability distribution, r : S → R is the reward function, ρ 0 : S → R is the distribution of the initial state s 0 , and γ ∈ (0, 1) is the discount factor.
Let π represents the stochastic policy of an agent that maps the action a t for a given s t to the future state s t+1 , and R(π) represents the expected discounted reward of the policy: , and s t+1 ∼ P(s t+1 |s t , a t ). The stateaction value function, Q π , is defined as [14]: An agent interacts with an environment by sequentially selecting actions, observing states, and receiving rewards to find the optimal policy that maximizes the cumulative discounted reward R t = ∞ l=0 γ l r t+l . Different algorithms have been proposed to approximate Q π (s t , a t ), e.g., Q-Learning, Deep Q-Learning [14,1]:

PROPOSED METHOD
In this section, we describe the proposed exploration strategy with Policy Augmentation. The goal of the proposed method is to augment the policy based on what the agent has experienced so far using IMC methods. Our assumption is that the agent has neither any previous experiences, nor the model of the environment. We describes the proposed IMC method for recovering the Q function in Section 3.1.
The goal of IMC methods is to recover the underlying low-rank matrix representation of the action-value function from the values of a small given subset of state-action pairs. If either state spaces or action spaces are continuous, we use hashing, e.g., SimHash function [7], to discretize their spaces and form the Q matrix.
Let Q ∈ R M ×N be the actual action-value matrix that the agent will learn, andQ ∈ R M ×N represents the recovered action-value function by the IMC method, where M is the total number of states and N is the total number of actions. Let Ω t represents the set of the observed entities of the action-value matrix-Q Ωt -at time step t, and Q Ωt (s, a) refers to the value of the state s and action a. After each time step, the agent receives a reward from the environment, R t+1 , that changes the value of either an observed entity of Q t or add a new member to the observed set Ω t .
IMC methods incorporate additional information about actions and states, making them appropriate for predicting the values of state-action pairs that have not been explored. Let X ∈ R M ×m and Y ∈ R N ×n represent the state and action features (additional information), respectively. The sth row of X represents the features of the sth state, and the ath row of Y is the features of the ath action. Algorithm 1 shows the proposed exploration strategy for training DRL algorithms. The strategy gets the DRL algorithm, state and action features (X and Y ), the total number of time steps (τ ), and initial state s 0 as its input. Because the goal is to generate high-value rollouts during the early episodes, we set a threshold τ e . The proposed algorithm uses the recovered augmentedQ to take actions while t < τ e . During this time period, DRL collects high-value rollouts for training. For t > τ e , the exploration and exploitation of RL Take a t , receive s t+1 and r t .

11
Collect rollouts for training the DRL algorithm. 12 if t < τ q then 13 UpdateQ t using Eq. 1 and Q t (s, a) = Q t (s, a)/N q t (s, a).
14 Initialize U t and V t . 15 repeat 16 Update U t using Eq.3.

18
until convergence criterion is met; 19 Augment the policyQ = X T U t V T t Y T . 20 Train the DRL algorithm using the rollouts. 21 end agents is similar to other previously proposed methods, e.g., PPO, TROPO, Curiosity-based, etc. Note that Policy Augmentation cannot replace a learned policy by DRL because it is only a prediction. It is used during the first τ e episodes.
The main assumption of matrix completion and IMC methods is that the underlying matrix-Q function-is low rank. Thus, we calculate the IMC method for the first τ q steps (τ q < τ e ) while the Q function is still low-rank. As the agent visits more state-action pairs, the rank of the Q increases. The computational complexity of the proposed method is associated with the complexity of the proposed IMC method.

Policy Augmentation
The proposed IMC method recovers the Q function from a given set of observed state-action pairs. We assume that Q t is a low-rank matrix (during the first τ q time steps) while the agent is in the early stages of exploration and the set Ω t is very small: |Ω t | M × N , where |Ω t | is the size of the set Ω t . The normalized state-action value matrix is written asQ t = XW t Y T , where W t is an unknown, low-rank matrix W t = U t V T t ∈ R m×n with the unknown rank r min(m, n). Our goal is to recoverQ t fromQ Ωt by solving min Ut,Vt [9]. There exists a unique solution to this optimization problem if at least nm log(nm) state-action values are explored [15], and U t and V t are incoherent [16] 1 . 1 Interested readers may refer to [15,16] for the proof.
The contribution of our proposed IMC method is that it racks the changing set Ω t as new pairs are revealed during the exploration. Moreover, the values of the observed set Ω t changes with time because the agent receives more rewards (either positive, negative, or neutral) from the environment for repeating the previously taken actions (Eq. 1). Following the definition in [17,9,18], we consider Q 1 ,Q 2 , . . . ,Q t as a set of linear subspaces on the Grassmann manifold, forming a geodesic.Q i is a point in the Grassmann manifold and is represented as a pair (U i , V i ) [16,9]. At each time step, only one element of Q t is changing. Thus, the linear subspaces ofQ t−1 andQ t remain close to each other. The distance between two subspaces is defined by the canonical distance (geodesic or arc-length distance) [16]. We use a set of principal angles {θ j } r j=1 , r = min(m, n) to quantify d(U t , U t−1 ) and d(V t , V t−1 ) as the projection distance: . The augmented action-value matrix,Q is recovered by minimizing the cost function J with respect to U t and V t [10]: By taking the partial derivative of the cost function J with respect to U t and V t and set them to zero [9], the closed form expressions are obtained. Then, the values of U t and V t are updated iteratively similar to [10]. The matrices U t and V t are initialized with random dense matrices. The algorithms updates U t and V t in Eq.3 and 4 alternatively until the convergence criteria meets.
where (i, j) indicates the element of the matrix in the ith row and the jth column. We selected λ 1 = λ 2 = 1 here.

RESULTS
The OpenAI Gym toolkit [20] and the Stable-baseline package [21] are used to evaluate the proposed algorithm. Two environments, Mountain Car and CartPole, are selected. Policy Augmentation is integrated with Proximal Policy Optimization (PPO) [11] (PPO+ Policy Augmentation), and Deep Q-Learning (DQN) [1] (DQN+ Policy Augmentation). The performances of PPO+Policy Augmentation and DQN+Policy Augmentation are compared with PPO [11], DQN [1], PPO+ Counting [7], and AttA2C (Att+ Curiosity) [22] in both environments. The experiment is repeated 10 times with 10 random initialization, and each time for 2, 500, 000 rollouts. The performances of different methods and exploration strategies are compared using their cumulative rewards. Because the  We used Car Position and Car Velocity as the columns of X (state features), and Y = [−10 + , 0 + , 10 + ] T . epsilon is a small non-zero value, e.g., 1, that prevents multiplying W by zero. Figure 3.b shows the cumulative rewards of different agents. The darker lines represent the mean of the cumulative rewards of 10 repetitions, and the shaded areas represent one standard deviation of the cumulative rewards over 10 repetitions. The PPO+ Policy Augmentation and the DQN+Policy Augmentation have the best and the second best performances, respectively. Other methods are not able to generate high-value rollouts during the first 100 episodes.
The experiments showed that the performance of the proposed Policy Augmentation exploration strategy depends on the first taken action. In the beginning, the Q matrix is zero and the agent chooses randomly an action from all zerovalued actions. The estimation ofQ values are affected by which action is taken in the first step. To achieve a superior performance, the proposed strategy observes the cumulative reward. If the cumulative rewards during the first few episodes are below a predefined value, Q is reset to zero and the agent takes a new action randomly.

CONCLUSION
Policy Augmentation predicts the values of unexplored stateaction pairs using IMC methods. By taking actions at early stages according to the predicted state-action values, agents collect rollouts with high values. Deep reinforcement learning algorithms trained with the high-value rollouts learn policies faster than existing exploration strategies. Policy Augmentation results in faster training of deep reinforcement learning algorithms, which makes deploying deep reinforcement learning algorithms in new environments more practical.
Future work will consider developing more computationally efficient IMC algorithms that update the previously augmented policy after each step from the new observed entity. The future work will also develop DRL algorithms that balance the number of high-value and low-value rollouts for training in environments with sparse rewards.