Paper: Reward-Consistent Dynamics Models are Strongly Generalizable for Offline Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
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 \(P_t\) as the generative models while the dynamics rewards \(D_t\) as the discriminators.

Author response:
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.

Reviewer comment (Theory):
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?

Author response:
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.

|                                                       | Average Model MAE |
| ----------------------------------------------------- | ----------------- |
| original model                                        | 0.226             |
| model with transition filter (ensemble discriminator) | 0.127             |
| model with transition filter (single discriminator)   | 0.141             |

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 eﬃcient reinforcement learning with linear function approximation.” COLT 2020.

Reviewer comment (Evaluation):
On page 27, last paragraph indicates that in the halfcheetah task, the dynamics reward appears to be less robust on OOD samples—is this not contradictory to what the main paper claims? Perhaps this has to do with the coverage of the data set—it appears that both `halfcheetah-medium-expert` and `halfcheetah-medium` are visibly suffering in OOD generalization, but they have less coverage on the state-action space compared to `random` and `replay`.

Author response:
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.

|                                                       | Average Model MAE |
| ----------------------------------------------------- | ----------------- |
| original model                                        | 0.226             |
| model with transition filter (ensemble discriminator) | 0.127             |
| model with transition filter (single discriminator)   | 0.141             |

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 eﬃcient reinforcement learning with linear function approximation.” COLT 2020.

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

Author response:
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.

Reviewer comment (Experiments):
Experimentally, does this correspond to 5000 as suggested by Table 4 in the appendix?

Author response:
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.

Reviewer comment (Experiments):
In experimental analysis, it appears the paper focuses mainly on locomotion tasks, and they seem to be mostly overlaps (i.e. NeoRL and D4RL share very similar environments). I am curious if this applies to other non-trivial environments that are supported by D4RL.

Author response:
To address this concern, we run experiments on a more challenging Adroit domain, which is used in existing works like EDAC [An et al., 2021] and MOBILE [Sun et al., 2023]. Due to time limit, we choose two tasks `pen-cloned-v1` and `pen-human-v1`, and run the method MOREC-MOBILE. The results are reported in the following table. We see that MOREC-MOBILE consistently outperforms MOPO and MOBILE. We are still running experiments on other tasks and will provide more results in the future revision.

|               | BC     | EDAC            | MOPO   | MOBILE | MOREC-MOBILE             |
| ------------- | ------ | --------------- | ------ | ------ | ------------------------ |
| pen-cloned-v1 | $38.3$ | $68.2$          | $54.6$ | $69.0$ | $\mathbf{74.3 \pm 14.4}$ |
| pen-human-v1  | $25.8$ | $\mathbf{52.1}$ | $10.7$ | $30.1$ | $37.3 \pm 18.8$          |

Reviewer comment (Theory):
Page 2, lines 3-4. It is unclear why the task reward function is known is reasonable given the explanation—it perhaps reinforces that reward function should not be known since it can be considered to be part of the dynamics model?

Author response:
We apologize for this misleading claim. The correct claim is that if the task reward function is unknown, we can easily resolve this issue by incorporating it into the dynamics model and learning this augmented dynamics model.

Reviewer comment (Writing):
Page 2, last third line of first paragraph. What does it mean by "static capabilities"? Does it mean the model is no longer trained after (i) before equation (1)?

Author response:
Yes, you are right. We will revise this claim to remove any confusion.

Reviewer comment (Writing):
Page 1, paragraph 2, line 3: "on the models" instead of "on models".

Author response:
Thanks for your kind suggestions. In line 3 in Algorithm 1, the notion should be the empty set. We have fixed these typos and throughly revised our paper.

Reviewer comment (Writing):
Page 1, paragraph 2, line 4: "the model errors" instead of "model errors".

Author response:
Thanks for your kind suggestions. In line 3 in Algorithm 1, the notion should be the empty set. We have fixed these typos and throughly revised our paper.

Reviewer comment (Writing):
In Algorithm 1, line 3: Is \(\phi\) the same as the policy \(\phi\)?

Author response:
Thanks for your kind suggestions. In line 3 in Algorithm 1, the notion should be the empty set. We have fixed these typos and throughly revised our paper.

Reviewer comment (Experiments):
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?

Author response:
**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.

|                                                       | Average Model MAE |
| ----------------------------------------------------- | ----------------- |
| original model                                        | 0.226             |
| model with transition filter (ensemble discriminator) | 0.127             |
| model with transition filter (single discriminator)   | 0.141             |

**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](https://anonymous.4open.science/r/iclr24-paper-7694-response-materials-C836/) 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.

Reviewer comment (Experiments):
How does this method compare to a simple filtering technique like selecting the transition with the lowest bellman error?

Author response:
**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.

|                     | MOREC-MOPO               | MOPO WITH BELLMAN ERROR FILTER |
| ------------------- | ------------------------ | ------------------------------ |
| halfcheetah-med-exp | $\mathbf{112.1 \pm 1.8}$ | $59.1 \pm 9.6$                 |
| halfcheetah-med-rep | $\mathbf{76.5 \pm 1.2}$  | $57.7 \pm 2.5$                 |
| halfcheetah-med     | $\mathbf{82.3 \pm 1.1}$  | $67.6 \pm 2.7$                 |
| halfcheetah-rnd     | $\mathbf{51.6 \pm 0.5}$  | $41.3 \pm 2.18$                |
| average             | $\mathbf{80.6}$          | $56.4$                         |

Reviewer comment (Theory):
When filtering, why construct a softmax distribution instead of directly taking the argmax action?

Author response:
**Question 3:** When filtering, why construct a softmax distribution instead of directly taking the argmax action?

**Answer Q3:** We appreciate your insightful question. In some tasks such as `halfcheetah`, we observe that incorporating more uncertainty into the model rollout process could help improve the generalization performance of the generated rollouts. A similar observation have been presented in the existing MBRL work [1]. Therefore, we choose to construct a softmax distribution.

To empirically verify this point, we conducted additional ablation experiments. In the `halfcheetah` tasks, based on MOREC-MOBILE, we employed two strategies for transition filtering: argmax and softmax. The results are listed in the following table. The results indicate that the average performance of the softmax strategy surpasses that of argmax.

|                     | SOFTMAX                  | ARGMAX                  |
| ------------------- | ------------------------ | ----------------------- |
| halfcheetah-med-exp | $\mathbf{110.9 \pm 0.6}$ | $110.8 \pm 1.3$         |
| halfcheetah-med-rep | $\mathbf{76.4 \pm 0.6}$  | $68.9 \pm 2.9$          |
| halfcheetah-med     | $\mathbf{82.1 \pm 1.2}$  | $77.1 \pm 3.0$          |
| halfcheetah-rnd     | $51.7 \pm 1.3$           | $\mathbf{53.2 \pm 1.4}$ |
| average             | $\mathbf{80.3}$          | 77.5                    |

Reference:

[1] Kurtland Chua, et al. “Deep reinforcement learning in a handful of trials using probabilistic dynamics models.” NeurIPS 2018.

Reviewer comment (Theory):
How does the reward dynamics relate to energy-based models? It seems quite similar to an energy function on (s, a, s').

Author response:
**Question 4:** How does the reward dynamics relate to energy-based models? It seems quite similar to an energy function on (s, a, s').

**Answer Q4:** The dynamics reward model has no connection to energy-based models which represent the density of the offline dataset. It evaluates a transition pair according to its faithfulness, disregard of its density in the data.

Reviewer comment (Writing):
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.

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
**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.

Reviewer comment (Presentation):
Figure 2 is very difficult to parse and how this leads to the conclusion that dynamics rewards is superior to dynamics model is not clear.

Author response:
**Answer C3:** Thank you for pointing out this issue. We will revise the paper to make Figure 2 easier to understand.

To understand why the dynamics reward is helpful, we can compare Figure 2(d) and (e). The dynamics rewards of the true transitions (e) are higher than the rewards of the model (d), in almost all of the region. This shows that the dynamics reward model can serve as a scoring function to evaluate the truthfulness of transitions, even in the OOD region.

More straightforwardly, in Figure 1 in the [anonymized link](https://anonymous.4open.science/r/iclr24-paper-7694-response-materials-C836/), we plot the transition error of the original model transition and the model transition filtered by using the dynamics reward. It can be observed that the low error region has been largely expanded, particularly into OOD region. This clearly demonstrates the superior performance of utilizing the dynamics reward for transition filtering.

Reviewer comment (Reproducibility):
The additional complexity of learning dynamics reward and implementation complexity might limit the wide adoption of the proposed method.

Author response:
**Answer C4:** Thanks for pointing out this issue. Here we record the time and memory complexities of MOPO and MOREC-MOPO in the `walker2d-medium-replay-v2` task. The memory and running time of MOREC-MOPO is around 1.5 times of MOPO. We believe that such a computation complexity is acceptable based on the notable performance boost achieved by MOREC. For the implementation complexity, we will open-source the code upon acceptance to facilitate the wide application of our method.

|            | RAM (MB) | TOTAL TRAINING TIME COST (HOUR)                              |
| ---------- | -------- | ------------------------------------------------------------ |
| MOPO       | 3132     | 11.6 (model learning + policy training)                      |
| MOREC-MOPO | 4715     | 17.7 (model learning + dynamics reward learning + policy training) |

Reviewer comment (Writing):
What does ``Accordingly, the behavior policy $\pi_{\beta}$ is regarded as the “dynamics model”` mean? It's difficult to parse and seems out-of-context as the behavior policy is not even defined yet before this line.

Author response:
**Answer Q1:** We are sorry for this misleading argument. The behavior policy $\pi^{\beta}$ is defined as the policy used to collect the offline dataset $\mathcal{D}$. Here we mean that when treating the environment transition as an agent, the behavior policy plays the role of the dynamics model in the view of this agent. We will clarify this point in the revised paper.

Reviewer comment (Theory):
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.

Author response:
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.

Reviewer comment (Novelty):
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?

Author response:
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.