Average-Reward Reinforcement Learning with Trust Region Methods

1 Introduction

The deep reinforcement learning (DRL) achieves prominent progress in many fields [Silver et al., 2016; Vinyals et al., 2019; Senior et al., 2020]. Most of the modern RL algorithms aim at maximizing E [∑∞t γ t rt], where γ denotes a discount factor strictly less than 1. The discount factor undermines the contribution of future rewards on the current value and ignores the long-run gains [Sutton and Barto, 2018], which means that the reward received after k steps in the future is only worth γ k-1 of it to the current state. The existence of discount factor simplifies the theoretical analysis by constructing the contraction mapping and offers a faster and more stable online learning by reducing estimation variances [Kakade, 2001;

∗Corresponding author: Li Xia, Jun Yang

Marbach and Tsitsiklis, 2003]. However, the discounted criterion is not aligned with the natural metric in many real scenarios, and introduces unnecessary bias of estimates despite lower variance [Thomas, 2014; Schulman et al., 2016]. In empirical implementations, the discount factor serves as a hyperparameter and the algorithm performance is very sensitive to it [Duan et al., 2016]. Thus, the decision maker should make a compromise between the foresight of the decision and the learning stability.

Different from the discounted criterion, the average criterion directly focuses on the long-run average performance. In the everlasting real-world problems, such as manufacturing, power systems and traffic controls, the average criterion is best suited [Mahadevan, 1996; Dewanto et al., 2020]. Optimizing the average criterion completely removes the discrepancy between the training performance and the evaluating metric due to discounting. Therefore, the average criterion is well explored in the classical MDP literature [Howard, 1960; Puterman, 1994; Cao, 2007]. However, for the data-driven implementation, the average reward RL algorithms is much less investigated compared with the discounted criterion. It is relatively hard to study the average RL algorithms, since many properties are lost when setting γ = 1 in the discounted framework. An obvious evidence is that the Bellman operator is no longer a contraction mapping in the average setting [Bertsekas and Tsitsiklis, 1995].

One of the key difficulties in the average reward RL is to guarantee the values estimated from the sample paths are bounded without discounting. Many previous works extended the standard Q-learning to the average setting, such as Rlearning [Schwartz, 1993], RVI-learning [Puterman, 1994; Bertsekas et al., 1995], and CSV-learning [Yang et al., 2016]. In the methods above, a relative value (estimated average reward, reference state value or a constant from prior knowledge) is subtracted from the target value, while they may still cause divergence and lead to unstable training [Dewanto et al., 2020]. Compared with the value-iteration based ones, the policy-iteration based methods, such as policy gradient [Sutton et al., 2000] and Actor-Critic [Konda and Tsitsiklis, 2000], are perhaps more suitable for average reward algorithms due to more stable training and potential for large scale problems. However, they still suffer the instability problem of value estimation in the average setting.

In this paper, we solve the average reward RL problem by introducing the trust region approach for the average criterion. Based on Perturbation Analysis (PA) theory, we overcome the analysis difficulty and develop a unified trust region theory with discounted and average criteria. We provide a novel performance bound which tightens the previous one [Schulman et al., 2015; Achiam et al., 2017] when γ → 1. The bound provides a direct evidence that average criterion is preferred than the discounted one for ergodic MDPs. From the view of implementation, we provide a practical algorithm called Average Policy Optimization (APO). We address the difficulty of estimating value with the average criterion and provide a new method named Average Value Constraint to solve the value drifting problem. To the best of our knowledge, APO is the first average DRL algorithm with neutral networks as approximators. Finally, our experiments on the continuous control benchmark MuJoCo show that APO can beat the discounted PPO in the metric of average reward, which further confirms the superiority of our approach.

2 Preliminaries

2.1 Average RewardMDPs

Consider an infinite-horizon ergodic Markov decision process (MDP), defined by the tuple 〈S, A, P, r, d0〉, where S denotes a finite set of states, A denotes a finite set of actions, P : S × A × S → [0, 1] denotes the transition probability, r : S × A → R denotes the reward function, d0 : S → [0, 1] denotes the distribution of the initial state s0. Let π : S × A → [0, 1] denote a stochastic policy, τ denote a trajectory (s0, a0, s1, · · · ), and τ ∼ π denote that the distribution of trajectories depends on π: s0 ∼ d0, at ∼ π(· | st), st+1 ∼ P(· | st, at).

With the assumption of ergodicity of MDPs, we aim at maximizing the long-run average performance under π:

We are interested in the steady-state distribution dπ, which is defined as follows:

It is remarkable that as the effect of the current action in the infinite horizon is considered, the long-run performance is independent of the initial state distribution. Let Vπ(s) denote the state-value function, Qπ(s, a) denote the action-value function, and Aπ(s, a) denote the advantage function. The value functions are defined by subtracting the long-run average reward ηπ from the current reward to make the functions bounded [Howard, 1960]:

2.2 Discounted RewardMDPs

With the discounted reward criterion, we focus on the recent rewards by introducing the discount factor γ ∈ (0, 1). We optimize the policy to maximize its normalized expected discounted return:

The discounted performance ηπ,γ can be viewed as the average reward under discounted steady-state distribution dπ,γ:

Different from the average criterion, the policy performance with the discounted criterion is dependent on d0. The discount factor makes a smooth interpolation between d0 and dπ. To see that, we have limγ→0 dπ,γ = d0 (the bandit case) and limγ→1 dπ,γ = dπ (the long-run average reward case).

In the modern DRL literature, the value functions are defined as following:

As the value function is a core concept in RL, people may expect to optimize the long-run performance by letting γ → 1. Unfortunately, many important properties are lost when γ → 1, e.g., value functions will be unbounded which devastate the algorithm performance. Another example is the trust region theory for policy optimization, which we will discuss next.

2.3 Trust RegionMethod

The trust region methods [Schulman et al., 2015; Schulman et al., 2017; Achiam et al., 2017] are a series of DRL algorithms updating policies with the following approximate policy iteration:

where ¯D(π ‖ πk) = Es∼dπk ,γ [D(π ‖ πk)[s]] with D denoting some distance measure and δ > 0 is the step size. The policy set {π | ¯D(π ‖ πk) ≤ δ} is called trust region, in which it is safe to evaluate policies with the samples collected with πk. In theoretical analysis, D is often chosen as the total variation divergence defined by DTV(π′ ‖ π)[s] = 12 ∑a∈A |π′ (a | s) - π(a | s)|. In practice, Kullback-Leibler divergence is preferred, which is defined as DKL(π′ ‖ π)[s] = ∑a∈A π′ (a | s) log π′ (a|s) π(a|s) . Since DTV(p ‖ q) ≤ √DKL(p ‖ q)/2, the theoretical analysis is aligned with the practical implementation.

The primary motivation of the trust region method is that optimizing the surrogate objective in the trust region guarantees monotonic performance improvements. When the distance is bounded by DTV, the solution of the problem in (7) has the lower performance bound [Achiam et al., 2017]

where ϵγ = maxs |Ea∼πk+1 [Aπk,γ(s, a)]|.

The advantage of the trust region method majorly comes from two aspects. In each iteration, the samples collected with πk are reused in evaluating any π in trust region, which significantly improves the data efficiency compared with the traditional policy gradient methods. The other hand, the performance lower bound avoids collapse after bad updates when neural networks are used as the policy approximator [Duan et al., 2016].

However, existing analysis of the lower performance bound for the trust region method is not applicable when γ → 1, which is easy to be verified by limγ→1 γ/(1 - γ) = ∞. As the larger discounted factor and average reward setting are important in real problems, more delicate theoretical analysis is needed to understand the behavior of trust region method when γ → 1 even γ = 1.

3 Unified Trust Region Theory with Discounted and Average Criteria

In this section, we propose a unified trust region theory based on the PA theory that extends the existing analysis of performance improvement bound from discounted criterion to average case. Moreover, we provide a novel policy performance bound for police updating. With the help of this bound, performance monotonicity of policy improvement is guaranteed when γ → 1. Typically, our result shows monotonic policy improvement is also guaranteed with the average reward criterion, simply by setting γ = 1.

3.1 General Formulation

We redefine the discounted value functions in (6) as follows:

Although this formulaic approach may be not familiar to DRL researchers, it is commonly used in the MDP literature [Cao, 2007]. The main benefit of this value definition is that it bridges the gap between the discounted and average criteria. With the new definition of value functions, we obtain that limγ→1 Vπ,γ = Vπ immediately.

Next, we present the performance difference formula for any policies π, π′ .

Lemma 1. For any policies π, π′ , the following equation holds:

This relationship has been shown in many previous studies of trust region method [Kakade and Langford, 2002; Schulman et al., 2015; Achiam et al., 2017]. Actually, (9) can be viewed as a rewriting of the performance difference formula in a sample path form. The performance difference formula is the key result of PA theory [Cao, 2007] and it quantifies the change of the performance of Markov systems corresponding to the change of policies. For the readers unfamiliar with PA, we attach a brief introduction in Appendix A.1.

Powered by the performance difference formula, we derive the average performance difference formula at once.

Corollary 1. For any policies π, π′ , the following equation holds:

3.2 Policy Improvement Bound

While the performance difference formula accurately describes the performance difference for two arbitrary policies, it is difficult to use the formula to develop practical algorithms directly. The bottleneck is that to evaluate the performance difference, we need the samples from dπ′ ,γ. That is selfcontradictory since no trajectory of the new policy is available unless the policy is updated. Thus, instead of optimizing the difference formula directly, the trust region method optimizes the following surrogate objective:

where dπ′ ,γ is replaced by dπ,γ in (9) and the probability of the sampled actions is corrected by importance sampling.

It is natural to query what is the difference between the surrogate objective and the real performance difference. To answer this question, we present the bound as follows.

Proposition 1. For any two stochastic policies π, π′ , the following bound holds:

The result tells us that if the discounted steady-state distribution of two policies is close enough and the advantage function is bounded, we are safe to use the surrogate objective for policy updating.

Next, we build up the connection between the distance of the discounted steady-state distributions and the distance of policies.

Proposition 2. For any two stochastic policies π, π′ , the following bound holds:

and κπ′ is the Kemeny's constant of the Markov chain induced by π′ .

The above conclusion is the key contribution of this paper, which reveals that the difference of discounted steady-state distribution is bounded with the distance of the policies multiplied with a constant factor ξγ. This factor ξγ is a time constant that reflects how policy changes affect long-term behaviors of Markov chains. When γ is not very large, the long-run effect of the current decision gradually decreases, allowing decision makers to focus on a few steps after the decision and ignore the impact afterwards. However, if γ is large enough, we cannot use discounts to suppress the long-term impact of the policy anymore. Fortunately, when letting γ approach 1, we find that it is still able to bound the long-run impact by introducing Kemeny's constant [Kemeny and Snell, 1960], which means the average time of returning to the steady-state distribution starting from an arbitrary state.

We further explore the theoretical result of ξγ by observing following phenomena.

Remark 1. For the average reward criterion, we have ξ = limγ→1 ξγ = κπ′ - 1.

Remark 2. The largest ξγ is not taken at γ = 1. In particular,

We should not intuitively believe that a smaller γ always leads to a tighter performance bound. When γ is larger than (14), the lower bound decreases as the γ increases and finally becomes κπ′ - 1 at γ = 1. It supports that when the problem focuses more on long-term performance, the average criterion is better than a large discount factor in the trust region method, since the former is able to obtain a more accurate estimation.

Combining the preceding two bounds together, we conclude the following theorem.

Theorem 1. For any two stochastic policies π, π′ , the following bound holds:

ηπ′ ,γ - ηπ,γ ≤ Lπ,γ(π′ ) + 2ϵγξγEs∼dπ,γ [DTV(π′ ‖ π)[s]]. In particular, the bounds hold with the average criterion: ηπ′ - ηπ ≥ Lπ(π′ ) - 2ϵξEs∼dπ [DTV(π′ ‖ π)[s]], ηπ′ - ηπ ≤ Lπ(π′ ) + 2ϵξEs∼dπ [DTV(π′ ‖ π)[s]].

The result clearly shows that the performance difference is bounded by three components: the difference of policies, the maximum change in the values, and a time constant depending on the problem. It not only facilitates the design of practical trust region algorithms, but also inspires us to further understand the nature of general RL algorithms.

In the end of this section, we give another perspective for why the average criterion is preferred to the discounted one in trust region method. It should be emphasized that the above analysis is based on the discounted steady-state distribution, from which it is not convenient to sample. In the practical implementation, most discounted algorithms ignore the difference in state distributions and directly sample from the steady-state distribution dπ, which leads to biased estimates for the policy gradients [Thomas, 2014]. However, with the average criterion, performance estimation is unbiased as sampling under dπ is theoretically-justified.

4 Average Policy Optimization

Based on the unified trust region theory, we extend the current discounted trust region algorithms to the average reward criterion. We address the value drifting problem raised in the average setting, and present a technique named Average Value Constraint for better value estimation. By approximating the policy and value function with neutral networks, we develop an algorithm based on the Actor-Critic framework, which is named Average Policy Optimization.

4.1 Value Estimation

In the most DRL algorithms, the value function is approximated by a value network Vφ(s) with φ as the parameters, which is updated by minimizing the following loss with stochastic gradient descent (SGD) method

where n indexes a batch of states sampled with π and ˆ V (sn) is the target value for sn.

Under the discounted criterion, ˆ V (sn) can be calculated directly using Monte Carlo (MC) method with ˆ VMC(sn) = ∑∞t=0 γ t rn+t, or using the value of next step to bootstrap (TD): ˆ Vφ(sn) = rn + γVφ(sn+1) [Sutton and Barto, 2018], where the latter is preferred in most cases for low-variance estimates. In the average setting, we adopt the update rule of R-learning [Schwartz, 1993] to evaluate the target value with ˆ Vφ(sn) = rn - ˆ η + Vφ(sn+1), where ˆ η is the estimate of ηπ. To smooth the value of ˆ η over batches, we update ˆ η by a moving average method: ˆ η ← (1 - α)ˆ η + α 1N ∑Nn=1 rn, where α is the step size.

Let φ′ denote the updated parameters:

where β is the learning rate. As the value network is updated by SGD, noise in the estimate is inevitable. Suppose that there is a bias ϵV between the value function and the true value Vφ(s) = V˜ φ (s) + ϵV , ∀s ∈ S. The real updated parameters based on the V˜ φ is

As β is sufficiently small, the post-update value function can be well-approximated by linearizing around φ using Taylor's expansion:

Subtracting the first equation from the second one, we have

For γ < 1, the bias between the real values and approximated values is gradually eliminated during training. However, this feedback mechanism disappears for γ = 1 as Vφ′ (s) - V˜ φ′(s) ≈ ϵV . That means the approximated values are drifting away from the true values during training due to noises. We call it the value drifting problem in the average reward setting.

The analysis above explains why the value estimation in average setting is inaccurate in the data-driven mode. Next we try to correct the estimation by taking advantage of the following property.

It tells us that the ideal estimates of values defined by (8) should have zero mean. We rewrite the problem in (15) by adding this condition as a constraint:

By constructing the Lagrangian function for the above problem, we have

where b := 1N ∑Nn=1 [Vφ(sn)] and ν is the Lagrangian multiplier. Taking the gradient of LV (φ, ν), we obtain

The problem in (19) is equivalent to the following problem

where ˜ V (sn) := ˆ V (sn) - νb. Compared with the discounted variant, we explicitly control the bias in value estimation. We name this method as Average Value Constraint.

4.2 Policy Optimization

Many algorithms [Schulman et al., 2015; Schulman et al., 2017; Wu et al., 2017] have been proposed to approximately solve the discounted problem in (7) with a policy network. Here we choose PPO [Schulman et al., 2017] as the base to develop our average reward algorithm, and point out that the policy updating is similar with other trust region methods. Instead of optimizing the original objective with the constraint, we optimize the surrogate loss as follows:

where θ is the parameters of the policy network, ω(θ) = πθ (at |st) π(at|st) is the importance sampling ratio and ˆ An is a shorthand of advantage estimation ˆ A(sn, an). To better balance the variance and bias of policy gradients, we further modify the

Algorithm 1 Average Policy Optimization

Input: α, β, λ, ν, ε, K, N, M

• 1: Initialize θ, φ randomly.
• 2: Set ˆ η = 0, b = 0.
• 3: for k = 1, 2, · · · ,K do
• 4: Run policy πθk for N times to collect {sn, an, r(sn, an), sn+1}, n = 1, . . . , N.
• 5: Update ˆ η ← (1 - α)ˆ η + α 1N ∑Nn=1 r(sn, an).
• 6: Update b ← (1-α)b +α 1N ∑Nn=1 Vφ(sn).
• 7: Compute δn with (22) at all timesteps.
• 8: Compute ˆ A(sn, an) with (23) at all timesteps.
• 9: Update the θ with Jπ(θ) in (21) for M epochs.
• 10: Update the φ with JV (φ) in (20) for M epochs.
• 11: end for

discounted GAE [Schulman et al., 2016] to an average variant. Define the average TD residual as

Then the advantage estimator is

where λ is the parameter to compromise between bias and variance.

5 Experiments

We conducted a series of experiments to evaluate APO, which are used to answer two questions:

• Compared with the discounted PPO, does APO have better performance in environments suitable with average criterion?
• What are the factors that affect the performance of APO? Specifically, does APO benefit from the Average Value Constraint?

We choose the continuous control benchmark MuJoCo [Todorov et al., 2012] with the OpenAI Gym [Brockman et al., 2016]. The MuJoCo tasks are designed to control the robot to finish certain tasks such as running as fast as possible without falling. It is natural to use the average criterion for some tasks in MuJoCo, such as HalfCheetch in which the robot keeps running and is rewarded by its speed. Moveover, we also evaluate APO in the other experiments with terminal states to examine the generalization ability of APO beyond the theoretical assumptions.

Both APO and PPO are implemented based on a modular RL package in PyTorch named rlpyt [Stooke and Abbeel, 2019]. While the performance of PPO is largely dependent on the code-optimizations [Andrychowicz et al., 2020], we do not consider any tricks apart from the vanilla PPO for controlling variables to justify true impact factors. For each task, we run the algorithm with 5 random seeds for 3 million steps and do the evaluation every 2000 steps. In the evaluation, we run 10 episodes without exploration by setting the standard deviation of policy as zero. All the hyperparameter combinations we consider are grid searched, which are showed in Appendix B. The computing infrastructure for running experiments is a server with 2 AMD EPYC 7702 64-Core Processor CPUs and 8 Nvidia GeForce RTX 2080 Ti GPUs. The time consumed of each experiment varies from 1.2h to 2h on average according to the complexity of the task.

5.1 Comparison to Discounted PPO

We compare our APO with the standard discounted PPO with the different discounted factors, and select the combination of hyperparameters which achieves the best average performance. The results in Figure 1 show that APO surpasses all the variants of PPO with different discount factors. We observe that the performances of PPO vary drastically with the change of γ. Meanwhile, The relationship between performance and γ variation is not monotonic. For example, γ = 0.99 achieves a sound result in Ant, while the performance with γ = 0.999 is terrible. This confirms that the performance with the average criterion cannot be optimized by letting γ approach 1 in the discounted framework.

The environments shown in Figure 1 are suitable for average setting, as they have no terminal states (Swimmer and HalfCheetah) or hard to fall even with a random policy (Ant). It is natural that APO beats the discounted PPO in these tasks. The audience may be curious about the performances of APO in other tasks, such as Hopper, Walker2d and Humanoid. We also evaluate APO in them and show the results in Appendix C. As we expected, APO achieves sound results with the competent PPO but does not beat the best in Hopper and Walker.

However, when we change the metric from average episode return to average reward, we find that APO is better than discounted PPO, especially in Humanoid. It shows that APO achieves higher speed but ignores the terminal reward as it focuses on the average performance under the steady-state distribution. We believe it is more suitable to model the safety requirements as constraints, which is easy to be extended from APO [Achiam et al., 2017].

5.2 Ablation Study on Average Value Constraint

We further evaluate the proposed method Average Value Constraint for the value estimation in average setting. We fix the penalty coefficient ν at different levels to see the changes of average values during training. The performances are improved by 66.14% with the Average Value Constraint. As shown in Figure 2, a larger ν produces a tighter constraint on the average value. Without the average value constraint (ν = 0), the average values fluctuate dramatically during training. Since only the relative value of different states is meaningful for the decision, the bias of average value is detrimental to the policy gradient estimation. As the average value constraint is derived from the Lagrangian method, we can also optimize ν adaptively during training [Stooke et al., 2020].

6 Conclusion

We pioneer the study of the trust region approach with the average-reward RL. Based on the PA theory, we first prove a lower performance bound the average criterion and tighten the previous discounted bound with a large discount factor. With the novel lower bound, the monotonic policy improvement of the average trust region method is guaranteed. Furthermore, we analyse the value estimation in the average setting and propose a new technique Average Value Constraint to stabilize the value estimation during data training. In empirical experiments, our proposed algorithm APO achieves better performance over most PPO variants with different discount factors in MuJoCo, which demonstrates the superiority of the average criterion in many engineering problems. In summary, this work fills the gap in the theory of DRL with the average criterion and provides a practical algorithm with a more appropriate criterion for many practical problems beyond the current benchmarks with the discounted criterion.