Reward-Consistent Dynamics Models are Strongly Generalizable for Offline Reinforcement Learning

We sincerely appreciate the time and effort you have invested in reviewing our paper, as well as your insightful and constructive feedback. Your comments have greatly assisted us in improving the quality of our work.

Comment 1: For the dynamics reward learning algorithm, why is GAIL better than GAN? It seems like the offline data set is fixed according to Algorithm 1, thus you can treat the dynamics models as the generative models while the dynamics rewards  as the discriminators.

Answer C2: We would like to clarify that we do not claim that GAIL is better than GAN. Actually, they are for different purposes but with the same learning objective. In GAN, the generator inputs a random variable and outputs a data (e.g., image), meanwhile in GAIL the generator is reinforcement learning. Here, learning the dynamics reward involves a reinforcement learning task in general, and we employ GAIL framework. Although GAN can be applied in the specific setting that horizon=1 in our experiment, we believe GAIL is suitable for general situations.

Comment 2: Proposition 1 is a result for tabular setting. Furthermore, this discriminator only differentiates properly the in-distribution transitions---how does this discriminator enable OOD generalization?

• On page 27, last paragraph indicates....

Answer C2: To address your first concern, we extend Proposition 1 to the linear function approximation setting [1]. In this setting, the dynamics models and discriminators are represented by linear functions with respect to the features. Compared to the tabular result, the difference is that the state-action space size on error terms is replaced by the feature dimension. We will involve this result in the revised paper.

For your second question, unfortunately, we do not have a theoretical answer to explain the OOD generalization ability of the dynamics reward. Nevertheless, we can find that the use of the ensemble of discriminators (in Eq.(3)) can be helpful. Specifically, in the following table, we list the average model MAE of the original dynamics model, the dynamics model filtered by the ensemble of discriminators, and the dynamics model filtered by a single discriminator (the last discriminator). It can be observed that the single discriminator performs worse than the ensemble of discriminators. However, we believe this is not the final answer. We will continue to study the theoretical properties of the dynamics reward in the future.

Finally, we would like to point out that the experimental results in Figure 12 on page 27 is not contradictory to our main claim. We do not claim that the dynamics reward is always robust to samples from any level of out-of-distribution. Essentially, the dynamics reward is robust to OOD samples within a region but its robustness also has limit. The results in Figure 12 exactly validate this claim. On one hand, for OOD samples with MAE less than 1.0, the dynamics reward exhibits the remarkable robustness. On the other hand, the dynamics reward cannot show unlimited robustness for extremely OOD samples with MAE around 10.0. From the information-theoretic perspective, we believe that no method can achieve unlimited robustness to samples from any level of OOD. It is a promising direction to further expand the robustness boundary of our dynamics reward learning method.

Reference:

[1] Chi Jin, et al. “Provably efficient reinforcement learning with linear function approximation.” COLT 2020.

Comment 3: The rate is $1/\sqrt{T}$, so how big $T$ should really be in practice? Proposition 1 somewhat suggests that should be almost covering the whole state-action space---seems like in that case we can just directly model the dynamics using a table?

• Experimentally, does this correspond to 5000 as suggested by Table 4 in the appendix?

Answer C3: We appreciate your question. In practice, we choose the number of iterations $T$ such that the losses of the discriminator and dynamics model have converged. In our experiments, this indeed corresponds to 5000 as shown in Table 4.

In the tabular setting, $T$ should cover the state-action space and the dynamics model is represented by a table. As mentioned in Answer C1, we can extend Proposition 1 to the linear function approximation setting. In this setting, the dynamics model is represented by a linear function with respect to the feature, and $T$ does not need to cover the whole state-action space and only depends on the feature dimension.

We thank the authors for the detailed response.

Regarding the answers:
Comment 1: Perhaps this is why I don't completely understand why this setting is GAIL but not GAN. Based on Algorithm 2 the objective appears to be sampling from the offline dataset $\mathcal{D}$, which is supposedly under the i.i.d. supervised learning setting, so how is this RL since $\mathcal{D}$ will not be modified based on $D$ nor $P$.

Comment 2:

• Regarding linear MDPs: I agree the referenced paper will likely give you a result that is dependent on $d$, though I would prefer if the paper can specifically mention this.
• Regarding OOD capability: I believe we agree that the experiment does not indicate that the method generalizes to extremely OOD samples. Consequently, this raises the question of whether the paper has over-claimed the novelty. In the paper title, the phrase "Strongly Generalizable" may be misleading---I was expecting this approach to also deal with extremely OOD samples. I suggest the paper to avoid using the word "strong" as it appears to be an absolute term, but instead using relative terms (e.g. Better, Stronger, More Generalizable, etc.)

Comment 3: This raises the question of scalability. If we start considering image-based tasks (e.g. Atari or robotics), then the number of models required for this approach becomes prohibitively expensive as each models can have millions, if not billions of parameters. I believe the paper should indicate (even better if address) this limitation.

Thank you for the other responses, I acknowledge that I have read them. Due to the above questions, I keep my current score.

We appreciate your time to review and provide positive feedback for our work.

Question 1: When training the dynamics reward, the (s, a) pair is always sampled from the dataset, whereas the (s') is either from the dataset or from the adversarial dynamics model. This makes the dynamics reward robust to OOD (s') but not (s, a). How does the model generalize to OOD (s, a)? Can you evaluate the dynamics model explicitly on OOD (s, a) pairs and see if the dynamics reward is still higher with valid transitions

Answer Q1: Thanks for your insightful questions!

How does the model generalize to OOD (s, a). Unfortunately, we do not have a theoretical answer to explain the OOD generalization ability of the dynamics reward. Nevertheless, we find that the use of the ensemble of discriminators (in Eq.(3)) can be helpful. Specifically, in the following table, we list the average model MAE of the original dynamics model, the dynamics model filtered by an ensemble of discriminators, and the dynamics model filtered by a single discriminator (the last discriminator). It can be observed that the single discriminator performs worse than the ensemble of discriminators. However, we believe this is not the final answer. We will continue to study the theoretical properties of the dynamics reward in the future.

Explicitly OOD evaluation. The performance of the dynamics reward in OOD region can be evaluated from Figure 2(d) and (e). The scores in Figure 2(e) are higher than Figure 2(d) in almost all the area, indicating that the dynamics reward gives higher scores on valid transitions in OOD region. Besides, the results of the D4RL task are shown in Figure 5. For (s, a) pairs with an extremely large MAE (1.0), we can view them as OOD pairs since the dynamics model has extremely poor predictions on such pairs. Figure 5(c) clearly shows that the dynamics rewards of true transitions are higher than that of model transitions on these OOD pairs.

Moreover, to evaluate the dynamics model explicitly on OOD region, we plot the figures in Figure 1 in the anonymized link to compare the original dynamics model and the dynamics model filtered by the dynamics reward. It can be observed that the low error region has been largely expanded into the OOD region.

Question 2: How does this method compare to a simple filtering technique like selecting the transition with the lowest bellman error?

Answer Q2: Thanks for raising this interesting point. The bellman error is affected by many factors. Firstly, the error of Q function significantly affects the bellman error. How to learn an accurate Q function is still an open problem. Even if we assume that an accurate Q function is available (in this case offline RL has been solved), the bellman error is still affected by the scale of reward and Q values at state-action pairs. Therefore, bellman error is not a proper metric for filtering transitions.

We have conducted comparison experiments on the halfcheetah task. For a transition $(s, a, r, s^\prime)$, we calculate the squared Bellman error of $(Q_{\psi} (s, a) - r - \gamma Q_{\psi} (s^\prime, a^\prime))^2$, where $Q_{\psi}$ is the current Q-function, $a^\prime \sim \pi_{\phi} (\cdot|s^\prime)$ and $\pi_{\phi}$ is the current policy. The results are shown in the following table. We observe that the performance of this baseline is significantly worse than MOREC. This demonstrates that the bellman error is not a good metric for transition filtering.

Thank you for taking the time to review our paper, and for your insightful comments.

Comment 1: A major contribution of this paper is the new concept of the "dynamics reward", and the algorithm used to learn the dynamics reward model (GAIL). However, it was not clear to me what this dynamics reward should represent. It would be helpful to describe the optimal closed-form solution (if there was one), or include some additional discussion on an intuitive interpretation of it.

Answer C1: Thanks for your valuable comment. We have no optimal closed-form solution due to the minimax formulation. But it can be intuitive to understand the dynamics reward as a scoring function for transitions. An ideal dynamics reward assigns higher scores for real transitions than any other ones. That is, given any state-action pair $(s,a)$ and the real next state $s^\star$, $r^D(s,a,s^\star) > r^D(s,a,s')$ for any unreal next state $s'$. Once such dynamics reward is obtained, it is straightforward to see that the transition function can be faithfully recovered by maximizing the dynamics reward, which is the motivation of this paper. We will involve the above discussion into the revised paper.

Comment 2: It seems to me the dynamics reward model kind of serves as a density model of the training data? If it is, then why not just learn a density model? If not, then what is the difference?

Answer C2: From Answer C1, the dynamics reward model has no connection to density. It evaluates a transition pair according to its faithfulness, disregard of its density in the data.

We notice that a reader may misunderstand the dynamics reward as a density model from Figure 2(d). In this figure, the high scoring region is close to the high density region. This observation is due to that supervised model learning makes less errors in the high density region, and coincidentally the dynamics reward correctly recognizes the errors of the model. It can be found from Figure 2(e) that, when given real transitions, the high scoring region is dissimilar to the high density region. We will include the above discussion into the revised paper.

---

Thanks for your valuable feedback. We hope that the above answers can address your concerns satisfactorily and improve the clarity of our major contribution. We are looking forward to your further response.

Thank you for taking the time to review our paper and providing us with your valuable feedback.

Comment 1: The interpretation of dynamics model as an agent and introducing the concept of reward adds unnecessary complexity to the method and makes it a bit difficult to easily understand the method. Simply formulating the main idea with adversarial generative training and using the score D as a reward could make the paper be more simple and easy to understand.

Answer C1: We would like to clarify that the interpretation of dynamics model as an agent has been proposed in several previous studies, including [1, 2, 3, 4]. Such perspective is widely recognized and can help understand the motivation of learning the dynamics reward by IRL. Meanwhile, the single score $D$ cannot be simply used as the dynamics reward as it is coupled with the (learning) transition behaviors. We use an ensemble of $D$, which is also motivated from the perspective of treating dynamics model as an agent. We will clarify this in the revised paper.

References:

[1] Venkatraman Arun, et al. ”Improving multistep prediction of learned time series models.” AAAI 2015.

[2] Yueh-Hua Wu, et al. “Model Imitation for Model-Based Reinforcement Learning.” ArXiv: 1909.11821.

[3] Tian Xu, et al. “Error bounds of imitating policies and environments.” NeurIPS 2020.

[4] Zifan Wu, et al. “Models as Agents: Optimizing Multi-Step Predictions of Interactive Local Models in Model-Based Multi-Agent Reinforcement Learning.” AAAI 2023.

Comment 2: The paper introduces multiple hyperparameters and did quite extensive hyperparameters search (e.g., temperature, penalty, and threshold, ..). Making sure that the baseline is fully tuned with the similar resource given to the proposed method could be important for a fair comparison.

Answer C2: Thanks for your suggestion. We would like to point that we did not use significant additional resources for hyperparameters search compared to the main baselines MOPO and MOBILE. Here we list the hyperparameter search spaces of MOPO, MOBILE and MOREC-MOPO/-MOBILE on D4RL tasks.

For MOPO and MOBILE, we use the implementation in [Sun et al., 2023], which reports the hyperparameter search strategy in Appendix C.2 and C.3 in [Sun et al., 2023].

1. For MOPO, [Sun et al., 2023] performs a grid search on three parameters: penalty coefficient $\beta$ with 3 choices, rollout horizon $h$ with 2 choices and uncertainty quantifiers with 3 choices. Therefore, the search space size in MOPO is 18.
2. For MOREC-MOPO, it introduces two additional hyperparameters: the temperature $\kappa$ and threshold $r^D_{\min}$ in transition filtering. The hyperparameter search strategy is detailed in Appendix D.2. However, we did not tune the hyperparameters $r^D_{\min}$ and $h$. For $r^{D}_{\min}$, based on the result shown in Figure 5(a), we choose a fixed value of 0.6 across all tasks. For $h$, we choose a fixed value of 100 across all tasks. Thus we only take a grid search on two parameters: penalty coefficient $\beta$ with 6 choices, temperature $\kappa$ with 2 choices. The total search space size in MOREC-MOPO is 12.
3. For MOBILE, [Sun et al., 2023] performs a grid search on two parameters: penalty coefficient $\beta$ with 4 choices, rollout horizon $h$ with 2 choices. Therefore, the search space size in MOBILE is 8.
4. For MOREC-MOBILE, its hyperparameters and search strategy are the same as MOREC-MOPO. Thus, the search space size in MOREC-MOBILE is 12.

In summary, MOPO and MOBILE are tuned with comparable resources to MOREC-MOPO/MOBILE for a fair comparison.