Reflective Policy Optimization

On-policy reinforcement learning (RL) aims to learn an optimal mapping from a sequence of states to actions based on rewarding criteria acquired through trajectories generated by interacting with the underlying environment. Proximal Policy Optimization (PPO) (Schulman et al., 2017) is one of the most typical algorithms in this category, owing to its simplicity and effectiveness. It has been successfully applied in various domains, including Atari games (Mnih et al., 2015), continuous control tasks (Dhariwal et al., 2017), and robot control (Lillicrap et al., 2016). However, existing algorithms optimize the policy based on a state-action pair and do not directly consider the impact of subsequent states and actions in the trajectory. This limitation inevitably gives rise to the sample inefficiency problem.

In prior studies (Mnih et al., 2015; van Hasselt et al., 2016; Schulman et al., 2015, 2017; Haarnoja et al., 2018; Silver et al., 2014; Fujimoto et al., 2018), the prevalent approach involves optimizing the policy using the value function of the current state. The value function potentially contains information about the subsequent data. However, a pertinent question emerges: Is optimizing a policy solely based on the value function the fastest (optimal) path to convergence? The answer is no, as this approach may overlook other crucial factors.

To illustrate this answer, consider an environment with a “cliff”. If an agent takes an action leading to falling off the cliff under a state, it must learn to avoid the action and the associated state. Returning to this state is perilous and could potentially trigger the same action. Therefore, the agent must actively avoid this state to enhance safety. The preceding action leading to this state must also be avoided, anticipating the possibility of re-entering that state again.

A similar scenario unfolds in a “treasure” environment, where an agent performs an action resulting in a large reward. Subsequent data imparts positive and negative insights into previous states and actions. Therefore, optimizing the previous action should incorporate information from subsequent state-action pairs, not relying solely on the value function. Intuitively, leveraging subsequent data directly can expedite algorithm convergence and enhance sample efficiency. Unfortunately, most existing algorithms lack this capability, which directly exploits the relationship between pairs of trajectory data for policy optimization.

To address the above issues with better sample efficiency, we introduce a new on-policy algorithm that optimizes the current policy by explicitly considering the relationship between the previous and subsequent state-action pairs in the sampled trajectories. Specifically, the proposed algorithm evaluates the current state-action pair and the impact of the subsequent pair of trajectories. It provides a more comprehensive perspective than traditional value function-based policy optimization. This approach enables the optimized policy to adjust its actions based on positive and negative information from subsequent states, thereby effectively reflecting the policy. We thus name the proposed algorithm as Reflective Policy Optimization (RPO).

The RPO algorithm, as proposed, directly focuses on policy optimization rather than solely on evaluating the value function. This distinction separates it from multi-step reinforcement learning methods (De Asis et al., 2018; Duan & Wainwright, 2023; Hernandez-Garcia & Sutton, 2019). Our proposed algorithm takes a direct approach by incorporating previous and subsequent trajectory information for policy optimization, establishing more clearly theoretical properties. We also derive a novel policy improvement lower bound, illustrating that, in addition to ensuring the desirable property of monotonic performance improvement, our method effectively reduces the solution space of the optimized policy, thus significantly accelerating the algorithm’s convergence procedure. Furthermore, our method improves sample efficiency.

We incorporate our proposed algorithm with the PPO’s clipping mechanism (Schulman et al., 2017) to provide a practical implementation. Following standard settings, we validate the effectiveness of our algorithm by utilizing an illustrative toy example, shedding light on the underlying working mechanism of RPO. Additionally, we showcase superior performance on widely recognized RL benchmarks, such as MuJoCo (Todorov et al., 2012) and Atari games (Brockman et al., 2016).

In this section, we establish a recurrence form by providing the equation relationships before and after the replacement of TRPO. Further, we reach general conclusions by extending TRPO with subsequent state-action pairs.

The proof of this lemma is given in Appendix A.4.

Note that from this lemma, the difference between the original formula and the replaced one is relevant to the normalized discount subsequent state-action visitation distribution $\rho^{\pi}(\cdot|s,a)$. By the boundary of the right-hand side of the equation, it is easy to obtain Theorem 1 of the paper (Schulman et al., 2015) and Theorem 1 of the paper (Achiam et al., 2017). From this lemma, we constructed a relationship between the current visitation distributions $(s,a)\sim\rho^{\pi}(\cdot)$ and the next $(s^{\prime},a^{\prime})\sim\rho^{\pi}(\cdot|s,a)$.

The proof of this theorem is given in Appendix A.5.

This theorem gives a general form for the difference between the policy performance of $\pi$ and $\hat{\pi}$ by finite sums. This equation accurately represents the general gap between the performance of $\pi$ and $\hat{\pi}$ from the trajectory-based. It portrays that subsequent state-action pairs can also directly impact optimizing the current policy. We refer to $\sum_{i=0}^{k-1}\alpha_{i}L_{i}(\pi,\hat{\pi})$ as the generalized surrogate objective function.

A slight problem may exist if the generalized surrogate objective function is directly optimized. Consider $L_{1}(\pi,\hat{\pi})$ in Eqn. (2) as an example. We consider this function without delving into the specific form of the parameters. When the environment is unknown, it can only be optimized by sampling. Considering a special case, the function $L_{1}(\pi,\hat{\pi})$ is optimized by using a sample $(s_{0},a_{0},s_{1},a_{1})$, i.e., $L_{1}(\pi,\hat{\pi})\approx(r_{0}-1)r_{1}A^{\hat{\pi}}(s_{1},a_{1})$. If $A^{\hat{\pi}}(s_{1},a_{1})<0$ and $r_{0}-1<0$, it follows that $(r_{0}-1)r_{1}A^{\hat{\pi}}(s_{1},a_{1})=[(r_{0}-1)A^{\hat{\pi}}(s_{1},a_{1})]r_{1}>0$. This implies an increase in the probability of $a_{1}$. However, when $A^{\hat{\pi}}(s_{1},a_{1})<0$, we should decrease the probability of $a_{1}$. It’s a contradiction. Thus, the term “1” in $r_{0}-1$ may be incorrectly misleading for policy optimization despite the soundness of the theory. This situation exists when the environment is unknown.

Next, we measure the gap between the policy performance $\eta(\pi)$ and $\sum_{i=0}^{k-1}\alpha_{i}L_{i}(\pi,\hat{\pi})$.

The proof of the theorem is given in the Appendix A.6.

Based on Theorem 3.2 and Corollary 3.3, a general lower bound exists for the policy performance of $\pi$. This theory makes good theoretical sense, which helps to understand the generalized surrogate function. For the case where $k=1$, the $l_{1}$ norm constraints are replaced by KL constraints. This outcome aligns with the lower bound observed in TRPO.

Theoretically, the previous section gave a lower bound for the policy performance of $\pi$. Although the generalized surrogate function incorporates the current and subsequent state-action pairs of the trajectory, the inclusion of the term ”1” in the function $L_{i}(\pi,\hat{\pi})$ introduces ambiguity regarding how the subsequent pairs influence the behavior of the policy at the current state, potentially yielding positive or adverse effects. In this section, we have made slight modifications to the generalized surrogate function $L_{i}(\pi,\hat{\pi})$ in Eqn. (2), defined

It is a natural modification that avoids the problems caused by the “1” item. The following theorem measures the relationship between the function $\hat{L}_{i}(\pi,\hat{\pi})$ and $\eta(\pi)$.

The proof of this theorem is given in Appendix A.8.

From Theorem 4.1, the first term of the generalized lower bound is referred to as the new generalized surrogate function, while the second term is known as the penalty term. It is worth noting that TRPO (Schulman et al., 2015) is a special case of the generalized lower bound for $k=1$. By optimizing the generalized lower bound, we can get a monotonically improving sequence of policies $\{\pi_{i}\}_{i=0}^{\infty}$, satisfy $\eta(\pi_{0})\leq\eta(\pi_{1})\leq\cdots$. Next, we intuitively analyze the new generalized surrogate function. The difference between the function $L_{i}(\pi,\hat{\pi})$ and $\hat{L}_{i}(\pi,\hat{\pi})$ is very small, involving the removal of the number 1 from the ratios’ product. However, their intended meanings are quite distinct. The function $\hat{L}_{i}(\pi,\hat{\pi})$ can directly utilize the information between the current and subsequent state-action pairs to optimize the current policy. Optimizing this function does not encounter the issues discussed in Section 3.

With $k=2$, we will explain the optimization procedure in detail. The function $\hat{L}_{1}(\pi,\hat{\pi})$ contains the ratio of the pair $(s,a)$ and $(s^{\prime},a^{\prime})$. If $A^{\hat{\pi}}(s^{\prime},a^{\prime})>0$, it indicates that the action $a^{\prime}$ is deemed favorable, and its probability will be increased through algorithm optimization. Simultaneously, the state $s^{\prime}$ is likely fine as well. To return to this state, we should increase the probability of the action $a$ under state $s$. In contrast, if $A^{\hat{\pi}}(s^{\prime},a^{\prime})<0$, similar results will be obtained. The agent can reflect on its current behavior based on subsequent information. For $\hat{L}_{0}(\pi,\hat{\pi})$, the action’s $a$ probability can be optimizing using the advantage function $A^{\hat{\pi}}(s,a)$. Therefore, optimizing the current action $a$ will be influenced by both the current and subsequent advantage functions $A^{\hat{\pi}}$, taking them into account.

In this way, the optimized policy will likely foster the agent’s reflection, and we observe that optimizing the generalized surrogate function lacks this ability. By utilizing the same trajectory, the agent can acquire more information. Figure 1 depicts an experiment conducted in the CliffWalking environment. The results in Figure 1 indicate that optimizing the new surrogate function reduces the number of falling off the Cliff and faster after reaching the goal $G$. In the experimental section, we explain this phenomenon in detail.

Theorem 4.1 demonstrates that the generalized lower bound is optimized for any $k$. As $k$ increases, the generalized lower bound is optimized using subsequent samples to learn implicit relationships of the current and subsequent states and actions data. However, is it suitable when $k$ takes a large value? The answer is no. Let’s look at the $\hat{L}_{k}(\pi,\hat{\pi})$ function individually. This objective function comprises the product of the $k$ ratios and an advantage function. If the ratio is too high, it encounters the problem of high variance (Munos et al., 2016), which, in turn, impacts the algorithm’s stability. Given this weakness, a very large value of $k$ may not be practical. In the experimental section 5.1, we discuss the values of $k$ and observe that as long as the agent leverages the relationship between previous and subsequent state-action pairs, it enables the agent to fall into the Cliff less often and to reach the goal $G$ faster. The experimental results show similarity whether $k=2$ or $k=3$. Therefore, the main part of the following discussion is framed in terms of $k=2$. The following theorem shows that the modified generalized surrogate function has another nice property except for the monotonicity.

The proof of the theorem is given in Appendix A.9.

Note that when the old and new policies do not change much, the set $\varPsi_{1}$ is a solution space of TRPO, and the set $\varPsi_{k}$ corresponds to the solution space of the $k$-th generalized lower bound. The theorem 4.2 shows that the scale of the solution space of the policy is reduced when $k=2$. Under certain conditions, the optimal policy $\pi^{\star}$ belongs to the set $\varPsi_{1}$ as well as to $\varPsi_{2}$. Reducing the solution space is possibly more efficient in finding a good policy, and therefore, it is intuitive that the algorithm’s convergence procedure can be accelerated. Furthermore, it can improve sample efficiency. Similarly, we can define the solution space of $k=3,4,\cdots$ and use the same way to get $\varPsi_{1}\supseteq\varPsi_{2}\supseteq\varPsi_{3}\supseteq\varPsi_{4}\supseteq\cdots$. Note that $\pi^{\star}$ is in those sets. It reveals the benefits of using current and subsequent states and actions of trajectory data to optimize the policy. This result provides a promising theoretical basis for our algorithm.

To verify the effectiveness of the proposed RPO algorithm, we utilize several continuous and discrete environments from the MuJoCo (Todorov et al., 2012) and Atari games in OpenAI Gym (Brockman et al., 2016) extensively adopted in previous works. We conducted experiments with the Mujoco environment based on code from (Queeney et al., 2021) and Atari games based on code from (Zhang, 2018).

This paper proposes a simple on-policy algorithm called Reflective Policy Optimization (RPO). This method aims to combine before and after state and action information of the trajectory data to optimize the current policy, thus allowing the agent to reflect on and modify the action of the current state to some extent. Furthermore, theoretical analyses show that our proposed method, in addition to satisfying the desirable property of the monotonic improvement of policy performance, can effectively reduce the optimized policy’s solution space, speeding up the algorithm’s convergence procedure. We verify the feasibility and effectiveness of the proposed method by a toy example and achieve better performance on RL benchmarks.

In future work, since RPO optimizes the policy by leveraging information from past and future pairs, it is straightforward to integrate it into other Actor-Critic algorithms or maximum entropy methods, further exploiting its advantages. We anticipate that RPO will play a good role in developing new reinforcement learning algorithms.

We thank the anonymous ICML reviewers for their insightful and constructive comments. Dr. Junliang Xing was partly supported by the Natural Science Foundation of China under Grant No. 62222606 and 62076238.

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.