ASOR: Anchor State Oriented Regularization for Policy Optimization under Dynamics Shift

Q4: Limited and questionable theoretical results

A: We appreciate that the reviewer takes the precious time to read the theoretical results and the proofs. The reviewer's concerns are discussed as follows.

• Relationship between Thm. 4.2 and 4.4: To analyze the KL-contrained optimization problem, Thm. 4.2 assumes that the learning policy $\hat\pi$ minimizes the KL divergence in Eq. (1) to be smaller than $\varepsilon$. But in practice, we can only minimize the divergence between the empirical version of distributions on finite training data. It remains unclear that given finite samples, how close the divergence between real distributions will be with high probabilities and how will such divergence influence the final performance lower bound. Thm. 4.4 gives such finite-sample analysis by leveraging the generalization ability of the network distance.
• Network distance and KL Divergence: The network distance has been proved to have good generalization property in finite sample analysis, where the intuitively preferred KL divergence fails to provide such property [Arora et al., 2017; Xu et al., 2020]. So it is natural to consider the network distance to derive the finite sample analysis of Thm. 4.4. This will surely introduce divergence from the practical algorithm, but we think it is a necessary trade-off to get meaningful theoretical results.
• Assumption from [Xu et al., 2020]: We actually emphasized the reliance on the linear structure of reward function by writing

if the reward function $r_{\hat\pi,T}(s)=\mathbb E_{a\sim\hat\pi,s'\sim T} r(s,a,s')$ lies in the linear span of $\mathcal P$

in Thm. 4.4. It is true that our paper borrow some proving techniques for finite sample analysis from Xu et al. (which has been clearly stated in the appendix), but these two papers work on completely different problems. Xu et al. focuses on Imitation Learning with Generative Adversarial Imitation Learning (GAIL), while our paper focuses on policy optimization under dynamics shift. To our best knowledge, among existing literature for policy optimization under dynamics shift, our paper is the first to discriminate empirical and real distributions and give finite sample analysis. So the insight here is that a PAC bound like Eq. (6) can be proved for policies regularized by the network distance. Future work that involves a similar procedure may also derive a PAC bound and compare with ours.

• Linear horizon dependency: To the best of our knowledge, Xu et al. [9] is the first to point out the advantage of linear horizon dependency over quadratic one and use it to explain GAIL's superior performance over behavior cloning. The idea behind this advantage is simple: with different horizon dependencies, the imitation error will accumulate at different rates along the trajectory. Before and after that, several papers proposed quadratic bounds when analyzing performance differences in face of dynamics or policy divergence. The coefficients of these bounds vary across different problem formulations. MBPO [7] consider model-based RL with learned dynamics model that has a gap of $\varepsilon_m$ with the ground truth dynamics, leading to the following bound with quadratic horizon dependency.

SLBO [8] also focuses on model-based RL and gives the following bound with quadratic horizon dependency:

where $\kappa=\gamma(1-\gamma)^{-1}$. EGPO [10] considers the shared autonomy of a learning controller and a expert controller, proposing that the value difference $\hat V$ can be bounded as follows:

where $K_\eta'$ describes the policy divergence between two controllers and its coefficient also has quadratic horizon dependency. Due to the notoriously high rate in error accumulation of quadratic horizon dependency, the aforementioned papers come up with special methods for error reduction. For example, MBPO propose to rollout with the learned dynamics model for no longer than K steps, rather than the whole trajectory. EGPO propose to involve human participants to select a few steps that the expert controller should take over.

In this paper, we formalize and utilize the similar structures of MDPs with dynamics shift and derive a new performance bound with linear dependency. Thanks to this good property, we are able to design policy regularization on the stationary state distribution, which is computed along whole trajectories. We hope such brief literature review could inform the reviewer of the importance of linear horizon dependency.

We thank the reviewer for the positive and constructive feedback. We are encouraged that the reviewer acknowledged the technical soundness and the extensibility of the proposed algorithm. With respect to weaknesses and questions, we provide our responses as follows.

Q1: Motivation of the Methodology

A: The reviewer raised concerns on whether inaccessible states and different optimal state distributions are detrimental to policy learning. One direct outcome of inaccessible states is that the optimal trajectory in one dynamics will no longer be optimal in the other. Imitating states in such a trajectory may lead to degraded overall performance.

In the lava world example with dynamics shift in Sec. 3.1 and Fig. 1, if the policy learns from the optimal state trajectory in the upper environment, it will take more steps and receive more negative rewards to reach the goal compared with the optimal policy. Another example is autonomous driving in different traffic densities, as illustrated in Appendix B and Fig. 6. The optimal state trajectory under low traffic densities will include many states with high speeds. But policies regularized by such state distribution will be more likely to crash under high traffic densities and will not be optimal.

With regard to empirical evidences, baseline algorithms BCO, SOIL, and SRPO are motivated by the idea of similar stationary state distributions, which all show inferior performances compared with AGSA. Among these, SRPO follows a similar pipeline of reward augmentation, but still performs worse than AGSA in all involved experiments. We hope such discussions on the motivation of the methodology can convince the reviewer on the importance of introducing anchor states and performing policy regularization upon them.

Q2: Limited scope of dynamics shift

A: The reviewer raised concerns on the limited scope of MDPs with similar structures and failing to incorporate shift in the reward function. We discuss these two concerns respectively and will add the discussion to the revision.

• Dynamics Shift: It is true that this paper relies on similar structures of MDPs. But without such similar structures, data from one MDP would never be informative for training policies in other MDPs. The only choice would be to train separate policies in each MDPs. One of the contributions of this paper is that we both motivate from (Sec. 3.2) and formally define (Def. 4.1) the similar structures, exploit them in policy optimization, and design a new training algorithm that improves data efficiency.

  In fact, existing literature more or less rely on assumptions on the MDP structures. Off-dynamics learning methods, such as DARC [1], DARA [2], and H2O [3], are also built upon HiP-MDP and further assume access to the data distribution of the test environments. Our method is free of such additional information while achieving superior performance. SRPO [4] and Imitation from Observation methods [5,6] has a shared research scope with our paper, but is built upon a stronger assumption on the MDPs.

  We would also like to mention that the setting of HiP-MDP is already general enough to handle a decent number of real-world senarios, such as autonomous driving with different road, tyre, weather, and traffic conditions, recommender systems with different news trend, user behaviors, and user preferences, video games with different map components, obstacle layouts, and so on. All of these tasks contain similar MDP structures that can be utilized in training.

• Reward Function: We would like to clarify that dealing with shifts in reward functions will be a complete different research problem. To our best knowledge, none of the existing literature for dynamics shift [1,2,3,4,5,6] considers potential reward shifts. To handle reward shifts, a typical approach will be meta-learning that aims at fast test-time adaptation with few-shot samples, which will be out of the scope of this paper.

Q3: Clarity and writing

A: The critical issue we focus is discussed in the first paragraph of the paper (Lines 31-41) and briefly stated in lines 40-41: how can RL policies be efficiently optimized using data collected under dynamics shift?. It is formalized in Sec. 2.1 (lines 93-94): Policy optimization under dynamics shift aims at finding the optimal policy that maximizes the expected return under all possible $\theta\in\Theta$. Such issue is important not only in literature, but also in many practical scenarios. There are quiet a number of related literature in this field as discussed in Sec 2.2 (Related Work). Real-world applications, such as autonomous driving, video games, robotics locomotion, and recommendation systems, may encounter dynamics shift in the training environment, which traditional RL algorithms struggle to deal with. So it is vital to design more robust and efficient training algorithms.

Q5: How to select $\rho$ in line 270?

A: $\rho_1$ and $\rho_2$ are selected through a few rounds of hyperparameter tuning. Their effects are examined in the ablation studies in Table 2, where ASOR shows robustness to hyperparameter changes.

Q6: Conclusion from the experiments

A: As stated in Sec. 5 (Lines 360-365), we conduct experiments to investigate whether ASOR can efficiently learn from data with dynamics shift and can be a general add-on module to different algorithms in diverse application scenarios. With empirical results in four different RL tasks, we show that ASOR can be effectively integrated with several state-of-the-art cross-domain policy transfer algorithms and substantially enhance their performance.

Q7: Novelty over a related work [4]

A: SRPO [4] focus on a similar setting with this paper, but is based on a strong assumption of universal identical state accessibility. We demonstrate that such an assumption will not hold in many tasks and a more delicate characterization of state accessibility will lead to better theoretical and empirical results.

References

[1] Off-Dynamics Reinforcement Learning: Training for Transfer with Domain Classifiers. ICLR 2021.

[2] DARA: Dynamics-Aware Reward Augmentation in Offline Reinforcement Learning. ICLR 2022.

[3] When to Trust Your Simulator: Dynamics-Aware Hybrid Offline-and-Online Reinforcement Learning. NeurIPS 2022.

[4] State Regularized Policy Optimization on Data with Dynamics Shift. NeurIPS 2023.

[5] State-only imitation with transition dynamics mismatch. ICLR 2020.

[6] Offline imitation learning with a misspecified simulator. NeurIPS 2020.

[7] When to Trust Your Model: Model-Based Policy Optimization. NeurIPS 2019.

[8] Algorithmic Framework for Model-based Deep Reinforcement Learning with Theoretical Guarantees. ICLR 2019.

[9] Error Bounds of Imitating Policies and Environments. NeurIPS 2020.

[10] Safe Driving via Expert Guided Policy Optimization. CoRL 2021.

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 motivation, the scope of dynamics shift, and the theoretical results. We would be grateful for the opportunity to provide any needed further feedback.

We thank the reviewer for the constructive feedback. We provide our responses to reviewer's concerns as follows.

Q1: The motivation of the paper is not entirely convincing.

A: We refer the reviewer to the general response for discussions on the motivation of anchor states. In general, the practical algorithm in Sec. 3.4 design an approximated approach of identifying anchor states in continuous state spaces, where the exact anchor state will not exist.

Q2: The empirical results show limited improvement over baseline methods

A: We would like to clarify that our AGSA algorithm is evaluated as an add-on module to existing algorithms. It requires minimal adjustment to the original training pipeline, which is the reward augmentation with two density ratios. The improvement over base algorithms is much higher than other algorithms as add-on modules. Take the performance in D4RL tasks as an example. The DARA algorithm only improves the performance of the MAPLE algorithm by an average of 2.7%. SRPO can improve by 11.1%, while our AGSA algorithm can improve by 25%. AGSA is also general enough to be applicable to five different training tasks, including Gridworld, D4RL, MuJoCo, MetaDrive, and a large-scale video game, achieving constant performance improvements. So it may be unfair to conclude that the empirical improvement is limited.

Q3: Practical strategy for dynamically identifying anchor states

A: Actually we have already employed a practical algorithm for identifying a "soft version" of anchor states in continuous state spaces. We refer the reviewer to the general response for detailed discussions. Four out of the five involved training environments have continuous state spaces. As discussed in Q2, the practical approximation shows decent performance improvements in these environments.

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 motivation and the empirical results. We would be grateful for the opportunity to provide any needed further feedback.

Q4: Arbitrary $T_0$ in the loss

A: The advantage of an arbitrary $T_0$ in the loss is that we do not have to care about the specific dynamics that the optimal policy is in. In the proposed AGSA algorithm, the optimal state distributions from any dynamics can contribute to policy learning in new dynamics. This greatly increases the data efficiency of AGSA as shown in the empirical results. In practice, we do not need to choose a specific $T_0$, since an arbitrary $T_0$ fits in the optimization objective Eq. (2). We just mix all training data collected from environments with different dynamics. According to Eq. (3), any states with high state value or frequently visited by the learning policy can be used to obtain the density ratio $\frac{d\_{T_0}^{\*,+}\left(s\right)}{d\_T^\pi(s)}$ and $\frac{d\_{T_0}^{\*,-}\left(s\right)}{d\_T^\pi(s)}$ in AGSA's reward augmentation pipeline.

With respect to the concern on anchor states in $T_0$, $T_0$ can indeed cover all anchor states across different MDPs. This is because in Def. 3.1, an anchor state $s$ should have non-zero visitation count in all dynamics, including an arbitrary dynamics $T_0$. Intuitively, the anchor states are the intersection of all optimal state trajectories under different dynamics and will therefore be a subset of the states reachable by $T_0$.

We thank the reviewer for the in-depth and constructive feedback. We are encouraged that the reviewer acknowledged the clear motivation, novel regularization method, and comprehensive experiments. With respect to weaknesses and questions, we provide our responses as follows.

Q1: Problem formulation

A: The reviewer's speculation on the focus is correct. In our HiP-MDP formulation in Section 2.1 (Lines 93-94), policy optimization under dynamics shift aims at finding the optimal policy that maximizes the expected return under all $\theta\in\Theta$, i.e., all possible dynamics in the HiP-MDP. In the revision we will make it a formal definition and define the training environments, which have hidden parameters $\theta_1,\theta_2,\cdots,\theta_n\in\Theta$. Compared with the formulation of off-dynamics RL, our HiP-MDP formulation does not assume any prior knowledge to target environments, which can be a natural fit to tasks with highly dynamic and non-stationary environments, such as robotics locomotion in open environments, autonomous driving, video games with evolving map components, and recommender systems.

Q2: More details in Section 3.3 and 3.4

A: We apologize for the obscure statements in Section 3.3 and 3.4. We give more clear demonstrations as follows which hopefully make our ideas easier to follow. They will also be added to the revision if the reviewer finds helpful.

• Connection with Eq. (3) and Eq. (2): Sorry for the redundant subscript t in the original Eq. (3). The density ratio $\frac{d_{T_0}^{\*,+}\left(s\right)}{d_T^\pi(s)}$ in Eq. (3) is the first reward augmentation term in Eq. (2). Eq. (3) breaks down the density ratio $\frac{d_{T_0}^{\*,+}\left(s\right)}{d_T^\pi(s)}$ in Eq. (2) into two density ratios that are easier to compute.
• Getting Eq. (3) with RND: According to Lemma 3.2 (Lines 209-213), the density ratio $\frac{d_{T_0}^{\pi,+}(s)}{d_T^\pi(s)}$ can be obtained by maximizing $\mathbb{E}\_P[f^{\prime}(\omega(s))]-\mathbb{E}\_Q[f^{\*}(f^{\prime}(\omega(s)))] (\*),$ where $f$ is a predefined function, $P$ and $Q$ are cumulative probability distributions of the density functions $d_{T_0}^{\pi,+}(\cdot)$ and $d_T^\pi(\cdot)$. Therefore, to obtain the density ratio, we only need to draw samples from $d_{T_0}^{\pi,+}(\cdot)$ and $d_T^\pi(\cdot)$, and then maximize Eq. (*). As $d_{T_0}^{\pi,+}(\cdot)$ is the distribution of anchor states visited by all optimal policies under different dynamics, states that are sampled from it should have higher visitation frequency. We therefore select RND as one of the visitation measure in large-scale tasks with continuous state spaces. Its loss is $L\_{\text{RND}}=\mathbb E\_{s\in\mathcal R}\|f\_0(s)-f\_\theta(s)\|\_2$, where $\mathcal R$ is the replay buffer, $f_0$ is a randomly initialized network, and $f_\theta$ is the learning network with parameter $\theta$. States with smaller RND loss are regarded to have higher visitation counts and more likely to be sampled from $d_{T_0}^{\pi,+}(\cdot)$.
• Details of ESCP: The Environment Sensitive Contextual Policy learning (ESCP) method trains an RNN-based context encoder $\phi(z_t|s_t,a_{t-1},z_{t-1})$, where $z_t$ and $z_{t-1}$ are latent variables that are trained to align with the environment parameter. By training policies $\pi(a|s_t,z_t)$ conditioned on latent variables, ESCP has the ability to adapt to environments with dynamics shift. Our AGSA method aims at efficiently exploiting data collected under dynamics shift. AGSA can work as an add-on module over ESCP and substantially enhances ESCP's performance. We will include more details of ESCP as well as its notations in the revision.

Q3: Existence of anchor states and getting anchor states during training

A: We refer the reviewer to the general response for discussions on anchor states. In general, the practical algorithm in Sec. 3.4 design an approximated approach of identifying anchor states in continuous state spaces, where the exact anchor state will not exist.

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 existence of anchor states, the arbitrary $T_0$ in the loss, and details in practical approximations. We would be grateful for the opportunity to provide any needed further feedback.

We thank the reviewer for the positive and constructive feedback. We are encouraged that the reviewer acknowledged the clear and effective motivation, the theoretical analysis, the compatibility with existing algorithms, and the empirical validations of AGSA. With respect to weaknesses and questions, we provide our responses as follows.

Q1: Writing and Background Accessibility

A: We thank the reviewer for the tolerance of our explanations in the current version. In general, Sec. 3.3 introduces how we approximate the density ratios $\frac{d_{T_0}^{\*,+}\left(s_t\right)}{d_T^\pi\left(s_t\right)}$ and $\frac{d_{T_0}^{\*,-}\left(s_t\right)}{d_T^\pi\left(s_t\right)}$ in practice. Lem. 3.2 basically turns the density ratios into the argmax of an optimization problem, which requires to sample from the distributions $d_{T_0}^{\*,+}(\cdot)$, $d_{T_0}^{\*,-}(\cdot)$, and $d_T^\pi\left(\cdot\right)$. The rest of Sec. 3.3 discusses how to approximately sample from these distributions when exact anchor state distribution is unavailable. We will add more explanations on Lem. 3.2 and the binary observation state $\mathcal O_t$ in the appendix of the revised paper.

Q2: Anchor State Identification

A: We refer the reviewer to the general response. We discuss how we design the approximation on anchor states in Sec. 3.3, where exact anchor state distributions are unavailable.

Q3: Incentivizing divergence on non-anchor states

A: We conduct ablation studies for encouraging divergence on non-anchor states in the large-scale video game environment. According to the results in the following table, AGSA without considering non-anchor states has lower overall performance than the original AGSA. This can be attributed to its high trapped rate and low policy entropy. The trapped rate indicates how frequent the agent gets stuck and stand still until the end of the episode. The policy entropy denotes the diversity of the agent behavior. These two metrics indicate a poor exploration ability without divergence maximization, which also demonstrates the effectiveness of encouraging divergence on non-anchor states.