State Regularized Policy Optimization on Data with Dynamics Shift

In many real-world scenarios, Reinforcement Learning (RL) algorithms are trained on data with dynamics shift, i.e., with different underlying environment dynamics. A majority of current methods address such issue by training context encoders to identify environment parameters. Data with dynamics shift are separated according to their environment parameters to train the corresponding policy. However, these methods can be sample inefficient as data are used \textit{ad hoc}, and policies trained for one dynamics cannot benefit from data collected in all other environments with different dynamics. In this paper, we find that in many environments with similar structures and different dynamics, optimal policies have similar stationary state distributions. We exploit such property and learn the stationary state distribution from data with dynamics shift for efficient data reuse. Such distribution is used to regularize the policy trained in a new environment, leading to the SRPO (\textbf{S}tate \textbf{R}egularized \textbf{P}olicy \textbf{O}ptimization) algorithm. To conduct theoretical analyses, the intuition of similar environment structures is characterized by the notion of homomorphous MDPs. We then demonstrate a lower-bound performance guarantee on policies regularized by the stationary state distribution. In practice, SRPO can be an add-on module to context-based algorithms in both online and offline RL settings. Experimental results show that SRPO can make several context-based algorithms far more data efficient and significantly improve their overall performance.

Paper

Similar papers

Peer review

Reviewer 6EPu7/10 · confidence 5/52023-06-17

Summary

The paper proposes a context-aware method called SRPO (State Regularized Policy Optimization) to address the issue of training RL algorithms on data with dynamics shift, i.e., with different underlying environment dynamics. The authors reveal and exploit the property of similar stationary state distributions in environments with similar structures but different dynamics. The learned stationary state distribution in SRPO is used to regularize policies trained in new environments, leading to more efficient data reuse and cross-domain ability. The paper also provides a lower-bound performance guarantee on policies regularized by the optimal state distribution in other dynamics. The experimental results show that SRPO can be a flexible add-on design that makes context-based RL algorithms far more data efficient and significantly improve their overall performance.

Strengths

+ The paper is well-structured and easy to follow, with relatively rigorous experimental validations and inspiring theoretical analyses. + The proposed SRPO algorithm is a context-aware method that tackles the challenge of training RL algorithms on data exhibiting diverse dynamics shifts, a very general problem in many real-world scenarios. + The authors have identified and effectively leveraged the observation that environments with similar structures but different dynamics possess similar optimal stationary state distributions. This discovery facilitates efficient data reuse and demonstrates improved performance in both online and offline settings, as evidenced by the experimental results. + The paper makes a valuable theoretical contribution by providing a lower-bound performance guarantee on policies regularized by the optimal state distribution in alternative dynamics.

Weaknesses

+ Please refer to the subsequent question section.

Questions

##### Method: + While the state regularized objective in SRPO is indeed well-motivated and illustrated by examples in the Inverted Pendulum environment, it would benefit from additional theoretical evidence to firmly establish the assertion that environments with similar structures but different dynamics inherently share similar optimal stationary state distributions. It should be noted that such a claim may not always hold true in some environments, and further theoretical investigation is expected to determine the generalizability of this finding. + Though the sample-based surrogate mechanism for state density ratio estimation is relatively well-founded, there is still some confusion regarding the inability of $D_{\text{fake}}$ with lower state values to capture the stationary state distribution $d_\pi$. It remains a challenge to handle $D_{\text{fake}}$ effectively, as it tends to underestimate the quality of the training policy and may even introduce a mixture of $d_\pi$ and $d_{\pi_\beta}$ induced by the behavior policy in static datasets in offline settings. ##### Experimental results: + DARC and DARA objective is derived from constrained optimization with dynamics regularization, instead of state regularization in SRPO, which both decouple the regularization term with policy that might vary in environments with dynamics discrepancies. It would be beneficial and interesting to provide more extensive analyses or explanations about these two types of regularization and demonstrate in what situations SRPO can perform larger advantage and why SRPO can work with more cases. It might also accout for the results of offline experiments that "MAPLE+DARA" offers very competing performance against "MAPLE+SRPO" on halfcheetah and walker tasks. + Further elaboration is needed to support the statement that SRPO enables efficient data reuse, specifically regarding reduced data consumption from source and target domains. + If there are any specific design considerations aimed at stabilizing the training process, it would be beneficial to include a detailed list of implementation details related to discriminator training. ##### Typos: - line 139 in Appendix: repeated "static".

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The authors have adequately and precisely addressed the limitations in this work.

Reviewer ZrWb6/10 · confidence 3/52023-07-06

Summary

The paper considers leveraging existing data in RL but with dynamics shifts. The proposed method augments the reward with a density ratio estimation term that encourage the policy to recover the optimal state distribution under the distribution shift. The theory motivates the algorithm design by showing that under certain assumptions, optimal policies under similar dynamics have close stationary state distribution. Finally, the experiments show that the proposed method indeed works well in practice in both online and offline setting with dynamics shift.

Strengths

1. The paper is well-written and easy to follow. 2. The proposed method can be easily added to some previous methods. 3. The paper provides motivating theory results to back up the design of the practical algorithm. 4. The design of the density ratio estimation design is well motivated in Section 3.3. 5. The paper provides good coverage of the experimental setting (both online and offline and includes several environments in the mujoco benchmark). The result is pretty good compared with previous baselines.

Weaknesses

1. The theory section (section 4) only provides some motivation for the algorithm design, instead of actually analyzing the algorithm. 2. The Lipchitzness of the reward function and dynamics still seems pretty strong even with section A.5: it should argument not only the Lipchitzness condition holds, and also the coefficients should be reasonably small.

Questions

N/A

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer PXH44/10 · confidence 3/52023-07-06

Summary

This paper presents an approach to RL/policy learning in the case where the dynamics parameters are potentially time-varying. The approach relies on the view that the optimal state distribution is similar across (small) variations in the underlying parameters, subject to assumptions. Thus, the authors propose an approach that constrains the distance between the realized state distribution under the trained policy and that of the optimal policy. Because the optimal policy state distribution is unknown, the authors propose to approximate the ratio of state distributions via discriminative methods, which in turn rely on an RL-as-inference formulation that leverages the value function.

Strengths

- The problem is reasonable, and guaranteeing (in some sense) reasonable RL performance with time-varying parameters is an important step toward practical robustness. - The method is reasonable and original. - The experimental evaluation is reasonably comprehensive.

Weaknesses

- The underlying motivation of similar state distributions is questionable. The similarity of these distributions is said to arise when the dynamics are similar (in this case, as reflected by (T, eps) notation) and when the reward functions and dynamics functions are Lipschitz. Then, the reward difference is approximately bounded by the products of the Lipschitz constants and the eps difference between dynamics functions. The authors claim these assumptions are mild---on the contrary, these conditions are quite strong and the resulting bounds are weak. In general, the idea of targeting the same state distribution doesn't seem reasonable for many real world systems, especially since a common use-case of RL is to handle non-smooth dynamics. - The approach is quite complex. The addition of a value-dependent discriminator is a substantial addition to standard RL workflows, and this added complexity will likely be a barrier to adoption. The authors should consider investigating simplified approaches---perhaps motivated by the same goals---that are more likely to see adoption. - Most importantly, the difference MAPLE + SRPO over MAPLE is extremely minor. Indeed, across many tasks the difference in performance is not statistically significant. In general, the combination of the high added complexity and the small resulting performance difference makes it difficult to recommend acceptance of the paper.

Questions

- Is it possible to avoid the training of new discriminator while still obtaining the benefits of state distribution regularization? - Is it possible to generalize analysis beyond smooth systems?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

The conditions stated in the paper (discussed above) are described as mild. The conditions are not mild, and the limitations of these assumptions should be discussed in more detail.

Authorsrebuttal2023-08-18

Replying to Reviewer PXH4

We sincerely value your dedicated guidance in helping us enhance our work. We are eager to ascertain whether our responses adequately address your primary concerns, particularly in relation to the underlying assumptions, the complexity of the approach, and the performance comparisons. We would be grateful for the opportunity to provide any needed further feedback.

Reviewer mfzT6/10 · confidence 4/52023-07-10

Summary

The paper aims to improve sample efficiency for RL algorithms trained in environments with dynamics shift. The core motivation is that, for different environments and their corresponding optimal policies, the state visitation distribution could be similar. Building upon this motivation, the authors propose a constrained policy optimization (CPO) that encourages the learned policy to have a similar stationary state distribution given by the optimal state distribution. For theoretical analysis, the authors propose homomorphous MDPs, which require two MDPs to reach the next state with the same support given the same current state and action. The result is that, given the reward and dynamics functions of MDPs having Lipschitz properties, if a policy $\hat{\pi}$ has a similar stationary state distribution with the optimal policy in one MDP $M$, $\hat{\pi}$ will have a lower-bound performance guarantee in all MDPs that are homomorphous with the MDP $M$. Experimental results show that SRPO outperforms existing state-of-the-art methods in both online and offline RL settings.

Strengths

The motivation that, for similar MDPs, their optimal policies should have similar stationary state distributions is very intuitive and makes a lot of sense. This motivation is well supported by both theoretical analysis and experimental results.

Weaknesses

Please see the questions section.

Questions

1. In equation 5, we use $R_\max$ directly. What if $R_\max$ is unknown? I assume we can estimate $R_\max$ simply by running several rollouts, or set $R_\max$ manually from domain knowledge. And what if the maximum reward we can obtain is different (or even has a different scale) for different states? E.g., we have $R_\max$ as 100, but for state $s_1$ the maximum reward is only 0, and for state $s_2$ the maximum reward is 100. Given equation 5, we always subtract $r(s_t, a_t)$ by 100 for both $s_1$ and $s_2$, but this seems unfair for $s_t$. 2. Line 159, there is no D.3 in the appendix. 3. In some cases, the assumption that the optimal policy for similar MDPs has similar stationary state distributions might not hold. For example , driving a car (with a fixed power level) on roads with different slopes, and the reward function is the velocity of the car. Apparently, the maximum velocity on different roads can be different, thus affecting the stationary state distribution. We can find more examples like this, only if the stationary state distribution for the optimal policy is strongly correlated with the MDP parameter. We do not expect an algorithm to be applicable in all cases, but I think we should clearly state the scope of this algorithm. 4. In both Figure 1 and the experiment section, we use a discrete set of MDPs. Several papers mentioned in line 104 "MDPs with Different Dynamics" indeed consider the case of continuous sets of MDPs, like VariBAD. For VariBAD, the generalization between different MDPs is carried out by the continuity of policy w.r.t. MDP dynamics. And it seems inappropriate to conclude "data with different dynamics are used in an ad hoc manner. The large amount of data with different dynamics are discarded during training, leading to poor sample efficiency (line 38)". Data collected from one MDP is helpful for other MDPs that are close to it. Can we 1) run experiments on continuous sets of MDPs, 2) make the discussion in line 37 more accurate? And if time permits, consider the VariBAD baseline.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors did not address limitations or potential negative societal impact of their work. No further improvement should be made.

Reviewer 6EPu2023-08-14

Thanks for the rebuttal. These supplementary explanations and experiments can definitely strengthen the authors' empirical claims in this paper, while I still have some questions and advice: > **"It would be beneficial and interesting to provide more extensive analyses or explanations about the two types of regularization (DARA style and SRPO style) and demonstrate in what situations SRPO can perform larger advantage and why SRPO can work with more cases."** I agree with the explanation regarding the relative ease of solving the HalfCheetah and Walker tasks, and their resilience to dynamics shifts. However, it would greatly enhance the paper if the authors could offer additional substantiation and analyses to elucidate why augmenting rewards with the **logarithm of stationary state distribution ratio** - as implemented in SRPO - outperforms the alternative approach of augmenting rewards with the **logarithm of dynamics ratio**, as observed in DARC-style algorithms such as DARA. Both lines of works hold theoretical validity, and further insights could provide a deeper understanding. My guessing could be that approximating the ratio of stationary state distributions lends itself to a more convergent and stable estimation process with neural networks. > **"As for online experiments, Fig. 4 of our paper demonstrates that in 7 out of 8 experiment settings, SRPO consistently outperforms baseline algorithms throughout the training process, even with very few training steps."** I see that SRPO doesn't show notably better performance in the initial training steps. In many experiments depicted in Figure 4, CaDM and CaDM+SRPO exhibit similar average returns in the first 2 million steps. The distinction becomes clearer when considering the long-term or asymptotic performance, where SRPO seems to have an advantage. > **Failure cases.** I recommend that the authors could add the potential failure cases to the appendix for further inspirations to this community since the theoretical assumptions might not always hold for complex real world cases though it can coverage most of the common ones.

Authorsrebuttal2023-08-15

Replying to Reviewer 6EPu

Thank you for your timely response and detailed questions/suggestions! We provide discussions on your concerns as follows. Q1: **Comparison between SRPO and DARC style of reward augmentation** A1: We concur with the idea that approximating the ratio of stationary state distributions can be more stable and converge quickly. The reason behind this is that in practice, DARA employs two distinct discriminators that estimates the ratio $\frac{q\left(\text {source} \mid \mathbf{s}, \mathbf{a}, \mathbf{s}'\right)}{q\left(\text{target} \mid \mathbf{s}, \mathbf{a}, \mathbf{s}'\right)}$ and $\frac{q(\text {source} \mid \mathbf{s}, \mathbf{a})}{q(\operatorname{target} \mid \mathbf{s}, \mathbf{a})}$ separately. The discriminators will receive $(s,a,s^\prime)$ and $(s,a)$ as input, respectively. In contrast, SRPO employs a single discriminator that estimates the ratio $\frac{p(\text{Real}|\textbf{s})}{p(\text{Fake}|\textbf{s})}$ . It solely requires the state $s$ as input. The additional inputs of action and next state can **increase the complexity of the discriminator in DARA and lead to a slower convergence**. Also, the joint distributions over $(s,a,s')$ and $(s,a)$, which appear in the denominator of the ratios, are more prone to have values close to zero. This can make the reward augmentation unstable and have high variance. Motivated by the intuition of similar optimal state distribution, SRPO focus on a more fundamental property in environments with dynamics shift and provides a simpler and more efficient approach of reward augmentation. This can explain why SRPO can outperform DARA in most of the offline RL tasks. To empirically substantiate the aforementioned discussions, we select the task of Walker-medium-expert and show the loss of the different discriminators of SRPO and DARA on the 100th epoch of training. Notably, the discriminator in SRPO demonstrates a smaller loss. The complete loss curve will be added in the revision. | Discriminator Input | $(s)$ (SRPO) | $(s, a)$ (DARA) | $\left(s, a, s^{\prime}\right)$ (DARA) | | :--- | :---: | :---: | :---: | | Loss on the 100th epoch | 0.84 | 1.03 | 1.12 | Q2: **SRPO doesn't show notably better performance in the initial training steps** A2: In online experiments, SRPO works as an add-on module to the CaDM algorithm. Apart from reward augmentation, all other modules and hyperparameters remain unchanged. Therefore, in the initial training steps the performance of SRPO+CaDM can be constrained by the exploration setup of the base algorithm CaDM, which is itself a value-based maximum-entropy reinforcement learning algorithm based on SAC. Since the training steps are the total interaction steps with all environments, in the first 2M steps the agent only interacts with each environment for 667k steps when there are 3 different dynamics, and for 400k steps with 5 different dynamics. The online policy can have high entropy and be highly stochastic with such few training steps. As a result, SRPO+CaDM may not show notably better performance in the initial training steps. Q3: **The theoretical assumptions might not always hold for complex real world cases though it can coverage most of the common ones.** A3: Potential failure cases will indeed be helpful for further exploration of the similar assumptions. We will add a failure example of maze environments with different obstacle layout, which violates the requirement of homomorphous MDPs (lines 223-256). We will also add an example task of stock market trading with very high transaction fees, which tend to have a high Lipstchitzs coefficient (lines 232-234). We will add the aforementioned discussions in the revision of our paper. We sincerely appreciate the questions and suggestions you provide in the two responses. We believe that solving the questions and adding the relevant discussions will greatly enhance our paper.

Reviewer 6EPu2023-08-15

Thanks for the reply. The additional explanations and analyses address my concerns properly and would be a strong addition to this paper, so I would keep voting for an acceptance. Wish you all the best with your publication.

Authorsrebuttal2023-08-15

Replying to Reviewer 6EPu

We thank the reviewer again for the valuable feedback. As we have addressed the concerns pointed out during the official review, would you like to kindly consider raising your score accordingly?

Reviewer mfzT2023-08-16

Thank you for submitting your rebuttal. The authors have made a commendable effort to address and answer my concerns, which is greatly appreciated. In relation to my third question, I kindly suggest that the authors consider incorporating additional intuitions into the main body of the paper rather than relying solely on theoretical statements around Line 242. Additionally, I would like to highlight that another reviewer has also raised the issue of failure cases, which I concur with. It would be highly beneficial if the authors could address this aspect as well.

Authorsrebuttal2023-08-16

Replying to Reviewer mfzT

Thank you for your positive feedback and recognition! We are grateful for your time and efforts in reviewing our paper and giving constructive comments. We will add the additional intuitions in the main body of our paper and add the failure cases in the appendix.

Authorsrebuttal2023-08-18

Replying to Reviewer mfzT

We thank the reviewer again for the valuable feedback. As we have addressed the concerns pointed out during the official review, especially the additional baselines, would you like to kindly consider raising your score accordingly?

Reviewer ZrWb2023-08-17

Response

I acknowledge the authors' rebuttal and I remain my rating towards acceptance.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC