Paper: A Unified Framework for Reinforcement Learning under Policy and Dynamic Shifts (ICLR-2024)

Reviewer comment (Theory):
I believe the proof of Proposition 3 could be clearer and more didactic with a better explanation on the assumptions and on why some steps are taken. Perhaps starting with an initial “rationale” behind the proof (describing the strategy to be followed) would be helpful. In fact, this could be also extended to the Section 4.2, which would help clarify some of the steps taken to arrive in the tractable learning objective (see my question below).

Author response:
Thanks for your suggestions. We apologize for the over brief discussion of theoretical derivation in Section 4.2, due to the length limits of the paper.

* To facilitate better understanding the proof of Proposition 3, we have included a sketch of the theoretical flow for the theory of Section 4.2 in Appendix A.
* We have provided the definition and the application of Fenchel Conjugate--the initial “rationale” of Proposition 3--in Appendix A.3 in our revision.
* Regarding the introduction of Bellman flow constraint:
  * Yes, we incorporate this constraint to ensure the tractability of optimization, a technique inspired by previous DICE works [1, 2, 3]. The Bellman flow constraint is essentially the definition of occupancy measure $d$, which enforces its behavior to satisfy its definition and thus will not change the original optimization problem. 
  * The key benefit of introducing this constraint is that, if we write the corresponding constraint problem into its Lagrangian dual form (Eq.(21) in the Appendix), and apply Fenchel conjugate transformation (Eqs. (19) and (22)), we can use the change of variable trick similar to DualDICE [4] and AlgaeDICE [5] (Eq.(23)) to get rid of the non-obtainable $\rho_{\widehat{T}}^{\pi}$ and convert the problem to a tractable form that only involves sampling from $\rho_{\widehat{T}}^{\widehat{\pi}}$ (Eq.(24)).

Reviewer comment (Experiments):
I believe the work provided sufficient empirical evidence to support the proposed framework. Nevertheless, one question remains: does OMPO scale for harder problems? For instance, does it work in meta-RL benchmarks such as Meta-World ML10, ML45? If not, why?

Author response:
Thanks for your suggestions. Meta-World ML10, and ML45 are meta-learning tasks, which involve multiple tasks with different rewards, such as button press and drawer open. Our proposed OMPO is a single-task RL algorithm rather than a multi-task algorithm, thus is not directly applicable to these tasks. Nevertheless, we think it is possible to use OMPO as a backbone algorithm to extend it to the multi-task/meta-learning setting. We will explore it in our future works.

Reviewer comment (Experiments):
The paper does not well describe the limitations of the proposed technique. It only states the challenges of tuning the buffer size but does not bring more practical information. For instance, I am curious to know about the stability of the method and how sensible it is for other learning parameters (such as those on Table 2). Additionally, it would be interesting to describe the computational resources needed to run the method for the presented benchmarks.

Author response:
* In our ablation, we have investigated the parameters, the local buffer size $\mathcal{D}_L$ and the weighted factor $\alpha$. 
* To further address the reviewer's concern, we have conducted ablations on the order $q$ of conjugate function in Figure 21, Appendix F.7 in our revision. The results show that $q\in[1.2,2]$ can yield satisfactory performance, with $q=1.5$ showing superior results in our experiments.
* Regarding the computation performance, although looks complex, OMPO actually does not have a high computational cost, its training time is at a similar level as SAC. We have provided computing infrastructure and training time in Appendix G in the revision.

Reviewer comment (Writing):
In Section 3, while describing the MDP, I believe the paper refers to the initial state distribution by two different symbols ($\mu_0$ and $\rho_0$). Based on the rest of the paper, I believe that the $\rho_0$ should be replaced with $\mu_0$.

Author response:
Thanks for your kindly reminder. We have corrected the typo, by replacing the $\rho_0$ as the $\mu_0$.

Reviewer comment (Reproducibility):
The pseudo-codes provided in Algorithm 1 should provide more training details. For example, calculating Eqs (23) and Eqs (25) requires sampling $s_0$ from the distribution $\mu_0$. How does the proposed method perform this operation exactly? I suggest the authors provide more training details, at least in the appendix.

Author response:
Thank you very much for this valuable suggestion! To deal with $s_0\sim\mu_0$ in Eqs (23) and Eqs (25), we adopt an initial-state buffer to store initial state $s_0$ similar to the implementation of previous DICE works [3,4]. When calculating Eqs (23) and Eqs (25), we can sample $s_0$ from the initial-state buffer. We have provided this detail in Appendix C in our revision.

Reviewer comment (Evaluation):
The proposed OMPO enjoys a low variance across different random seeds in terms of performance given stationary environments, as shown in Figure 2. Can the author provide some insights into this phenomenon?

Author response:
* OMPO outperforms the baselines under the Stationary environments because it can more effectively handle policy shifts. Note that in stationary environments where $\widehat{T}=T$, objective (5) reduces to:


$\widehat{\mathcal{J}}(\pi)=\mathbb{E}\_{(s,a,s')\sim\rho^\pi_T}\left[\log r(s,a)-\alpha\log(\rho^\pi_T/\rho^\widehat{\pi}_T)\right]-\alpha D_f\left(\rho^\pi_T\Vert\rho^\widehat{\pi}_T\right)$


$=\mathbb{E}\_{(s,a)\sim\rho^\pi}\left[\log r(s,a)-\alpha\log(\rho^\pi/\rho^\widehat{\pi})\right]-\alpha D_f\left(\rho^\pi\Vert\rho^\widehat{\pi}\right)$


$=\mathbb{E}\_{(s,a)\sim\rho^\pi}[\log r(s,a)]-\alpha\left[D_{KL}(\rho^\pi\Vert\rho^\widehat{\pi})+D_f(\rho^\pi\Vert\rho^\widehat{\pi})\right]$

which essentially regularizes the discrepancy between on-policy occupancy $\rho^{\pi}$ and the occupancy induced by off-policy samples $\rho^\widehat{\pi}$ from the replay buffer, which helps to alleviate potential instability caused by off-policy learning [1,2]. We have provided a more detailed discussion about the performance improvement in Appendix B in our revised paper.

* Yes, we also used the double-Q technique for OMPO, same as TD3 and SAC.

Reviewer comment (Theory):
Eqs (23) and Eqs (25) minimize the $Q(s, a)$ specifically for $s\sim\mu_0$. Since $Q$ is parameterized by a neural network, the $Q(s, a), s\sim\mu_0$ can be minimized spuriously low. How do the authors combat this potential training stability?

Author response:
In our empirical experiments, we did not observe any numerical issue with minimizing $Q$ with $s\in \mu_0$. It should be noted that in many existing DICE-based offline RL algorithms [3, 4, 5, 6, 7], the same $\min \mathbb{E}_{s\sim \mu_0} Q(s,a)$ term is also used, and no notable computation issue is reported. We think in the online setting, the possible instability issue is even less severe than in the offline case, as the online interactive environment will randomly reset after each episode ends, potentially providing an even broader distribution for $\mu_0$.

Reviewer comment (Reproducibility):
The proposed implementation is complicated. It introduces modules such as estimating the density ratio $\rho_T^\pi\left(s, a, s^{\prime}\right) / \rho_{\widehat{T}}^{\widehat{\pi}}\left(s, a, s^{\prime}\right)$ and performing min-max optimization. This can make the training unstable.

Author response:
- First, the instantiation of OMPO actually is not very complicated, the only additions in OMPO as compared to typical actor-critic frameworks are a local replay buffer and a distribution discriminator. 
- Second, learning an extra discriminator for policy learning is actually quite common in methods like GAIL [1] as well as lots of DICE-based RL/IL methods [2,3,4,5]. To estimate the density ratio $\rho_T^\pi\left(s, a, s^{\prime}\right) / \rho_{\widehat{T}}^{\widehat{\pi}}\left(s, a, s^{\prime}\right)$ stably, we extend the implementation of GAIL [1] from $\rho(s,a)$ to $\rho(s,a,s')$, whose training stability has been well verified.
- Empirically, we find the benefit of correcting policy/dynamics-shift using the extra discriminator significantly outweighs the possible instability in the discriminator. As shown in our main results in the paper, OMPO is stable and obtains superior performance, and in most cases, achieves even lower variance during policy learning.

Reviewer comment (Experiments):
The experiments are not thorough enough. See Q3 and Q4.

Author response:
* We thank the reviewer's valuable suggestion. In our original paper, we set the environmental parameters comparable to those in the baseline papers ($0.5\sim2$ times the original parameters). Besides, our considered dynamics shifts encompass multifaceted variations, such as simultaneous changes in wind and gravity.
* To fully address the reviewer's concern, we have conducted additional experiments in Non-stationary environments with a broader range of $0.5\sim3$ times gravity changes. We have included the results in Figures 19 and 20, Appendix F.6 in our revision. The results show, even with a more severe shift in dynamics, OMPO can still achieve much better performance than the baselines. Here are the brief results on the average return of the Ant and Humanoid tasks in non-stationary dynamics.

|                                  | Ant    | Humanoid |
| -------------------------------- | ------ | -------- |
| OMPO ($0.5\sim2$ times gravity)  | $2133$ | $3185$   |
| CEMRL ($0.5\sim2$ times gravity) | $1278$ | $1546$   |
| OMPO ($0.5\sim3$ times gravity)  | $1950$ | $2848$   |
| CEMRL ($0.5\sim3$ times gravity) | $1265$ | $1256$   |


Actually, all our methods more or less involve policy shifts, since our method as well as most baselines are off-policy algorithms, the training data sampled from the replay buffer always have some gaps with the current on-policy distribution. In Figure 2, we specifically compared our method on stationary environments (no dynamics shift) against other mainstream off-policy RL algorithms like SAC, TD3 as well as the DICE-based method AlgeaDICE, and OMPO achieves consistently better performance due to better handling of the policy shifts. In Figure 7, we also visualized the distribution discrepancy caused by policy shifts in different training stages, which shows that even in the stationary environment, the impacts of policy shifts can be clearly observed.

Reviewer comment (Theory):
It is completely not clear why optimising the proposed objective (5) should yield good performance on a new unseen environment, or even on the same environment off policy. In fact, there are several logical errors in the arguments.

Author response:
We think there are some misunderstandings here. As clearly illustrated in Figure 1, the non-stationary dynamics considered in our paper involve a series of different dynamics $\widehat{T}_1,\widehat{T}_2,...,\widehat{T}_h$ at different training stages, each representing a specific environment dynamics. This is a common setting considered by previous non-stationary RL studies [7, 8, 9]. We use $\widehat{T}$ to represent the transition dynamics that we encountered in the samples from the replay buffer of the **current training stage**. $\widehat{T}$ actually correspond to $\widehat{T}_1,\widehat{T}_2,...,\widehat{T}_h$ during the course of training. $\widehat{T}$ **does not** correspond to a stationary env, and we use the symbol of $\widehat{T}$ simply to keep the notations uncluttered and also provide a unified view of RL problem under policy/dynamics shifts.

Reviewer comment (Theory):
The authors propose to construct the occupancy measure $\rho_{T}^{\pi}(s,a,s')$ and assume that observations data has such distribution. However, stationary measures are generally ill-defined for non-stationary environments and it is not clear what this means. Thus when the authors write $\rho_{\hat{T}}^{\pi}(s,a,s')$, it appears that they must assume the data generating environment $\hat{T}$ is a regular environment, contradicting the premise of the paper.

Author response:
* First, we would like to politely remind that, $\rho_{\widehat{T}}^{\pi}(s,a,s')$ actually does not need to be explicitly computed in OMPO. If the reviewer checks Step 3 of Section 4.2, as well as its proof in Appendix A.3, by leveraging Fenchel conjugate and the change of variable trick similar to DualDICE [10] and AlgaeDICE [11], we can avoid evaluating $\rho_{\widehat{T}}^{\pi}(s,a,s')$ but solve a tractable formulation Eq. (12) that only involve samples from $\rho_{\widehat{T}}^{\widehat{\pi}}$ (i.e., the off-policy/off-dynamics samples from the replay buffer).
* Second, as we have discussed in the response to W1, using off-dynamics data to facilitate policy learning in the target domain is widely studied and acknowledged in literature [3, 4]. If the reviewer still thinks using off-dynamics samples is not relevant to policy learning in a new environment, we are happy to provide a comprehensive list of related references.
* Lastly, as demonstrated in our empirical results in Section 5.1, OMPO achieves strong performance under environments with shifted dynamics and substantially outperforms algorithms that are specifically designed for such settings, even in some highly non-stationary environments.

Reviewer comment (Evaluation):
It is generally not clear how specifically the introduction of $s'$ helps performance on the target environment.

Author response:
- As discussed above, our proposed objective (5) extends the general RL objective under policy & dynamics shifts, and the subsequent derivations are built upon the well-established DICE-related theories. Specifically, general RL objective $\arg\max\_{\pi}\mathbb{E}\_{(s,a)\sim\rho_\pi}[r(s,a)]$ has no treatment for policy and dynamics shifts, thus will be heavily impacted if encountering scenarios involving such shifts. In contrast, OMPO can handle these shifts carefully from our theoretical derivation using regularization terms $\log\left(\rho^\pi_T/\rho^{\widehat{\pi}}\_{\widehat{T}}\right)$ and $D_f(\left(\rho^\pi\_T\|\rho^{\widehat{\pi}}\_{\widehat{T}}\right))$, resulting in improved performance.
- We have added a logical flow for the theoretical derivation of our method in Appendix A of our revision paper, as well as extra discussions to explain the source of performance improvement in Appendix B. Please let us know if you have any specific comments on our theoretical derivation.