Policy Regularization on Globally Accessible States in Cross-Dynamics Reinforcement Learning

We thank the reviewer for the constructive and insightful comments. Responses are provided as follows. The linked file is available here.

Q1: Possibly empty accessible state set

A: Tasks with no globally accessible states are indeed extreme cases where ASOR may degrade into the base RL algorithm. But two practical factors can help alleviate such concern.

• One may regard the state dimension that remains distinct across dynamics as hidden environment parameter and do not include such information when identifying globally accessible states. For example, the robot height in the mentioned example can be excluded in the state space. States that are visited by expert policies with different robot height, e.g., smooth walking with moderate speed, can be regarded as globally accessible and preferred when training new policies.
• In the practical algorithm procedure of ASOR, a relative accessibility detection mechanism is applied. States that are most likely to be accessible in a training batch are regarded as preferable (Line 7 in Algorithm 1). We can rely on the discriminator network to automatically find appropriate states to assign higher rewards (Line 8 in Algorithm 1).

Meanwhile, the mentioned robotics scenario with different robot leg conditions can also be challenging for base RL algorithms which ASOR is built on, such as cross-domain RL methods. A considerable amount of training data may be needed to train decent policies in such scenarios.

Q2: Overlooked IfO papers

A: Thanks for bringing these papers into our attention. We will add discussions in the related work of the revised paper. We also agree with the reviewer that we should give IfO papers [1,2] enough credit for also considering state distribution mismatch. IfO methods are helpful when enough near-expert trajectories or unlabeled data are available and have better training efficiency. The advantage of ASOR lies in that its reward augmentation mechanism can be effectively integrated with both offline and online RL methods and adds minimal changes to their original training pipeline.

Q3: Additional baselines

A: Thanks for introducing alternative ways of constructing baselines. We refer the reviewer to Tab. 3 of the linked page, where comparative analysis on SOIL with 10% data and IQL is included. SOIL indeed benefits from the 10% data fix and shows increased performance, while IQL has similar performance with CQL. Both approaches cannot outperform MAPLE+ASOR. This can be because the offline datasets do not have enough state coverage for SOIL to imitate. Meanwhile, IQL and CQL still focus on state-action joint distributions in the offline dataset. Such distribution can be unreliable under dynamics shift since given the same state, the optimal action may be different in different dynamics. We are unable to reimplement the mentioned DICE-related methods due to the limited rebuttal time and will add them in the revision.

Q4: Potentially fragile training stability

A: Both the RND module and the GAN module are essentially supervised learning on given datasets. Their training will be more stable than the actor-critic module that involves RL loss. The red loss curve in Fig. 3 (right) also indicates that the RND loss and the GAN discriminator loss drop smoothly, indicating a relatively stable training process. One might relate the GAN module to the unstable training of GANs. But in fact, the GAN module only utilizes a GAN-like objective function. Its training procedure is different from GAN in that the data to train the discriminator is directly constructed from the replay buffer, instead of being adversarially generated.

Q5: Contradicting augmented reward

A: Fig. 3 (left) demonstrates the augmented reward of a single step, while in Fig. 3 (right) we show rewards averaged in a batch and multiplied by the coefficient. The extreme values of the augmented reward only exist when the walker agent is about to fall down near the end of the trajectory, so they only make up a small fraction of a training batch. The average augmented reward will therefore remain relatively stable during training.

Q6: Repeated $\lambda$

A: Thanks for pointing out this issue. We will change the Lipschitz coefficient to $\mu$ in the revision. The Lagrange multiplier $\lambda$ is the the same with the coefficient for the augmented reward according to Eq. (5), so we remain the same notation. For the values of $\lambda$ in different experiments, we refer the reviewer to Tab.1 of the linked pdf page.

Q7: Number of GPUs

A: We only use one V100 GPU for centralized training in Ray. Distributed environment rollouts are carried out with ~300 CPU threads.

We thank the reviewer for the constructive and insightful comments. Responses are provided as follows. The linked file is available here.

Q1: Robustness to hyperparams

A: Thanks for mentioning the hyperparameter table which is indeed missing in the original paper. We refer the reviewer to Tab. 1 and Tab. 2 in the linked file for detailed hyperparameters. In offline RL, we tried two values of $\lambda$: 0.1 and 0.3. For datasets with higher level of optimality, $\lambda=0.1$ is preferred so that less policy regularization is exerted. $\lambda=0.3$ is preferred on datasets with higher randomness to add more policy regularization based on optimal accessible state distributions. $\rho_1$ and $\rho_2$ can be similarly chosen based on offline data optimality with possible values (0.5, 0.5) and (0.3, 0.3). In medium-expert datasets, states are more likely to be sampled from the optimal accessible state distribution, so higher values $\rho_1$, $\rho_2$ are chosen. In online RL, $\lambda$ is set to 0.03 in MuJoCo environments and 0.1 in the Fall Guys-like environment, mainly to fit the reward scale of different RL environments. $\rho_1$, $\rho_2$ are kept to 0.5 in online RL experiments.

In Tab. 2, a fixed $\lambda=0.1$ has poor performance compared with MAPLE. This is because the augmented reward adds additional variance to environment reward. In some datasets its scale may not be large enough to exert reasonable policy regularization, so the reward variance dominates the training process and leads to a performance drop. Meanwhile, random partition denotes that the real data and fake data to train the GAN-like discriminator in Algorithm 1 are randomly chosen from the training batch. Its poor performance compared with MAPLE demonstrates the importance of constructing discriminator training data according to the accessible state distributions.

In Tab. 4 of the linked file, we conduct additional hyperparameter tuning where $\lambda$ ranges from 0.1 to 0.4 and $\rho_1,\rho_2$ equal to 0.3, 0.5, and 0.7. The results demonstrate that ASOR's performance is relatively stable with $\lambda=0.2, 0.3, 0.4$. $\rho_1,\rho_2=0.7$ may include too many states in the real dataset for discriminator training, so the performance drops accordingly.

Q2: Marginal improvements in HalfCheetah

A: We agree with the reviewer that ASOR is most beneficial in environments where accessibility varies, such as the fall-guys like game environment. As discussed in Sec. 4.2 (Lines 380~406), the agent will not fall over in the HalfCheetah environment, so the state accessibility will be more likely to remain unchanged under dynamics shift. This undermines the effectiveness of ASOR and leads to the marginal performance improvements.

Q3: More training steps in MetaDrive

A: We clip at 1M episodes to make training steps consistent with other environments. We will add the full training log in the revision.

Q4: The scale of fall guys-like game environment

A: The fall guys-like environment has a 1065-dimensional state space containing terrain, map, target, goal, agent, and destination information. Although there are no visual inputs, the state space contains diversified information that is challenging to integrate. Meanwhile, due to the highly dynamic environment and complex interactions between the agent and the environment components, a large number of environment steps is needed to train a decent policy. That is why we call the fall guys-like environment "large scale". The scale of the action space will not influence the effectiveness of ASOR, as the proposed policy regularization objective relies solely on the state distribution and do not take the action or policy distribution into account.

Q5: Restrictiveness of the assumption

A: As discussed after Def. 3.3 (Lines 208~219), the assumption is weaker than those in existing methods and less restrictive. It requires state $s'$ can be accessed from $s$ after dynamics shift, as long as it is accessible from $s$ in the original dynamics. It does not require a specific policy to do so and does not limit the number of intermediate states. Such assumption still holds in two of the mentioned extreme cases in that the state accessibility does not change.

Q6: Different optimal policy

A: According to Eq. (2,3), regularized policy optimization achieves performance lower-bounds without assumptions on how close the optimal policies are. Intuitively, the regularization provides a reasonable starting point for policies to explore more efficiently, but the starting point is not necessarily optimal. The walker2d environment are known to have many different near-optimal policies, while our ASOR algorithm can still achieve the best performance (Fig. 2).

Q7: Stability of GAN training

A: Due to the limited word count, we refer the reviewer to Q4 in the rebuttal to Reviewer UQdV.

We thank the reviewer for the constructive and insightful comments. Responses are provided as follows.

Q1: Determining accessibility in practice

A: Tasks with potentially empty globally accessible state set are indeed extreme cases where ASOR may degrade into the base RL algorithm. An example proposed by Reviewer UQdV is when we want to learn a walking policy for robots with either normal legs or crippled legs, and with different heights. In this case, as the robots may never perfectly achieve the same status, the globally accessible state will become an empty set. But two practical factors can help alleviate such concern.

• One may regard the state dimension that remains distinct across dynamics as hidden environment parameter and do not include such information when identifying globally accessible states. For example, the robot height in the mentioned example can be excluded in the state space. States that are visited by expert policies with different robot height, e.g., smooth walking with moderate speed, can be regarded as globally accessible and preferred when training new policies.
• In the practical algorithm procedure of ASOR, a relative accessibility detection mechanism is applied. States that are most likely to be accessible in a training batch are regarded as preferable (Line 7 in Algorithm 1). Even though there are no states that are perfectly accessible across dynamics, we can rely on the discriminator network to automatically find the most appropriate states to assign higher rewards (Line 8 in Algorithm 1).

Q2: Sampling from the accessible distribution

A: According to Sec. 3.5 of the paper, instead of directly sampling from the intractable accessible state distribution $d_{T_0}^{\*,+}(s)$ , we estimate the its likelihood ratio with $d_T^\pi(s)$. The key property used here is Prop. 3.8, which shows that the likelihood ratio can be obtained by optimizing a classifier similar to the GAN discriminator, as long as the real data are sampled from the state distribution in the molecular of the ratio and the fake data from the denominator. According to Eq. (4), the likelihood ratio $\frac{d_{T_0}^{\*,+}(s)}{d_T^\pi(s)}$ can be decomposed into the production of state optimality and accessibility. We therefore use the state value function and state visitation pseudo-count approaches to split the training batch and create classifier training data. In this way, we manage to obtain an estimation of the likelihood ratio, which serves as the augmented reward in practice.

We thank the reviewer for the constructive and insightful comments. Responses are provided as follows. The linked file is available here.

Q1: Practical estimation concerns

A: We discuss in Sec. 3.5 on how to obtain practical estimations. Instead of directly estimating the intractable accessible state distribution $d\_{T_0}^{\*,+}(s)$ , we estimate the its likelihood ratio with $d_T^\pi(s)$. The key property used here is Prop. 3.8, which shows that the likelihood ratio can be obtained by optimizing a classifier similar to the GAN discriminator, as long as the real data are sampled from the state distribution in the molecular of the ratio and the fake data from the denominator. According to Eq. (4), the likelihood ratio $\frac{d_{T_0}^{\*,+}(s)}{d_T^\pi(s)}$ can be decomposed into the production of state optimality and accessibility. We therefore use the state value function and state visitation pseudo-count approaches to split the training batch and create classifier training data. In this way, we manage to obtain an estimation of the likelihood ratio, which serves as the augmented reward in practice.

For the amount of transitions required, we refer the reviewer to Fig. (3)(right), where the red line (extra loss) contains the loss of training the estimation network. It takes roughly 1M steps to obtain a decent estimation, which is 1/6 of the total policy training steps. In other words, a relatively accurate likelihood ratio estimation is easier to obtain compared with a good RL policy, and can provide reasonable guidance in the policy training process.

When there are safety issues, safe RL techniques can be integrated as ASOR does not change the training pipeline of the base algorithms. One may also increase the coefficient of the augmented reward because it encourages the policy to stick to states that are more likely to be visited by expert policies in different dynamics.

Q2: Dynamics identification

A: Different base algorithms have their own approach of identifying environment dynamics, and ASOR as a general reward augmentation algorithm is agnostic to these approaches. For example, MAPLE and ESCP use context encoders trained with auxiliary loss to detect dynamics change. In the fall-guys like game environment, we include transformers in the policy network to automatically determine environment dynamics from history state-actions. In the lava scenario, we add environment information in the state space for simplicity. Apart from location coordinates, the state space has a 0-1 variable indicating whether there is a lava block nearby (Lines 153-158). In Fig. 1 (up), the agent at (2,3) will have (2,3,1) as the state; In Fig. 1 (bottom), the state will be (2,3,0). The agent can therefore be informed which situation it is facing.

We will add the aforementioned discussion to the revision.

Q3: Confusing notations

A: Sorry for the confusing notations. $d_T^*(s)$ is the brief version of $d_T^{\pi^*}(s)$ and means the same distribution. For the accessible state distribution such as $d^{\pi,+}_T(s)$, we plan to use a new distribution notation $\delta^\pi_T(s)$ for better readability.

Q4: Experimental designs

A: Thanks for recommending diffusion policies as baseline. We consider the Diffusion-QL algorithm and refer the reviewer to Tab. 3 of the linked pdf for comparative analysis. It exhibits better performance than CQL especially in the Hopper environment, but cannot outperform the average performance of MAPLE+ASOR. This can be because Diffusion-QL tries to recover the state-action joint distribution in the offline dataset. Such distribution will be unreliable under dynamics shift since given the same state, the optimal action may be different in different dynamics. For maximum entropy policies, we included SAC as baseline in online RL experiments. Meanwhile, ESCP and ESCP+ASOR are built upon SAC themselves and will benefit from the maximum entropy training objective.

For experiments with deviated dynamics during training and testing, i.e., testing with out-of-distribution dynamics, we refer the reviewer to Fig. 1 of the linked pdf. We utilize the OOD setting in ESCP, where test environment parameters, including the damping coefficient and the wind speed, are sampled from distributions that are 33% broader than the training distribution. In this scenario, HalfCheetah and Ant witness performance drops across different algorithms. Algorithms show almost no performance changes in the Walker2d environment on unseen test environments. This may be because dynamics shift exerts less influence on the walker agent. In both scenarios, ESCP+ASOR still achieves the best results, demonstrating its ability to efficiently train a more robust policy.

Q5: Literature completeness

A: Thanks for bringing LMDP and diffusion RL research into our attention. We will add discussions on these related works.