Paper: Efficient Offline Reinforcement Learning: The Critic is Critical (ICLR-2024)

Reviewer comment (Presentation):
First of all, the overall structure and writing of this paper necessitate meticulous reorganization and refinement. Some paragraph is confusing and hard to follow due to poor organization. Especially in Section 4 and Section 5, the important conclusion in these paragraphs need to be highlighted and summarized. In Section 5, the transition from Monte Carlo (MC) pretraining to emphasizing both actor and critic regularization is perplexing, especially since these regularization are not introduced in the methods section. And the title "Application to Adroit Environments" is incongruous as the methodology differs from the prior parts.

Author response:
We have restructured and re-written Sections 4 and 5 following your advice to better separate our experimental setup and implementation details from the results and discussion in each section. Each paragraph now corresponds to a separate point of analysis, and we have highlighted the key conclusions in each section. We have also renamed Section 5 from ‘Application to Adroit Environments’ to ‘Extension to Data-Limited Adroit Environments’, and added a Section ‘Motivation for Actor and Critic Regularisation with Pre-Training’ to explain the necessary extension of our methodology to handle the limited-data setting considered in this section.

Reviewer comment (Presentation):
The two parts of pretraining and regularization that the authors want to underscore appear to be incremental additions rather than naturally integrated components. This disjointed presentation detracts from the coherence of the paper and needs to be addressed.

Author response:
As mentioned above, we have now added an additional Section (5.1) to better justify our addition of combined actor and critic regularisation with pre-training. To summarise here, pre-training trains the actor and critic networks to first match the behaviour policy distributions, which improves the efficiency of subsequent off-policy RL, as we show in Section 4. However, in data-limited domains, subsequent off-policy RL can lead to policy collapse, where the performance drops to close to zero after pre-training. To prevent this, and to try to maintain at least the pre-training performance, we can add additional regularisation to both the actor and the critic to prevent chain extrapolation away from the initialised behaviour distributions. Therefore, our proposal is to pre-train both the actor and critic to match the behaviour policy, and then regularise them to prevent extrapolation significantly out-of-distribution of the behaviour distributions.

While we provide this additional justification of additional regularisation in Section 5.1, we also anticipate and highlight this idea throughout the paper, with the addition of regularisation being mentioned in the last paragraph of the introduction (Section 1), the overview of our approach in Section 3 (including Algorithms 1 and 2), and the abstract and conclusion.

Reviewer comment (Novelty):
Regarding the methodology, the paper's primary emphasis appears to be on the use of Monte Carlo (MC) estimates as pretraining targets. However, the results in Appendix A.3 indicate that the efficiency and performance during pretraining stem from the Behavioral Cloning (BC) loss rather than MC. MC only contributes to stability during subsequent fine-tuning. Furthermore, in Section 5, the authors assert that pretraining is less critical than both regularization techniques. Consequently, I am unconvinced about the significance of this work.

Author response:
Appendix A.3 provides an ablation which isolates the significance of the MC pre-training, so we are confused about why the reviewer is unconvinced of the significance based on these results. 

In Figure 2, we see that pre-training both the actor with imitation learning and the critic with a Monte-Carlo value error (blue) provides an efficiency benefit over randomly initialising these networks (black, still using additional LayerNorm). However, one might wonder if pre-training the actor is enough, since this is what leads to non-zero performance during pre-training (and does not depend on the critic values at all, since it is only trying to imitate and not optimise behaviour during this stage). 

Therefore, in Appendix A.3, we perform actor only pre-training with imitation learning. We see that during pre-training performance is equivalent as expected, but after pre-training performance drops off and takes a long time to recover, generally performing similarly to training from scratch with TD-learning as in Figure 2. Therefore, since the blue line reaches maximal performance more quickly than both the black line in Figure 2 and the red line in Figure 4 (Appendix A.3) we have isolated that the significant efficiency gain is due to our proposed MC pre-training.

In Section 5, we agree that pre-training does not help in these environments (as we mention) since performance collapses with or without pre-training. Therefore it is the combined regularisation approach that we propose that is responsible for the performance gains over the baselines in these data-limited environments. However, the pre-training stage is still beneficial for obtaining a non-zero initial baseline performance that allows for easier tuning of the regularisation parameters to prevent subsequent performance collapse.

Reviewer comment (Experiments):
The experiments are also limited in variety of dataset types and domains. For instance, BC pretraining may depend on data quality; therefore, additional dataset types such as "medium-replay" and "random" datasets are necessary to substantiate the importance of this work. Moreover, further ablation studies are required to validate the paper's claims. Current results in Figures 3, 5, and 6, which only base on a single dataset, are unconvincing. Additionally, the inclusion of more domains, such as AntMaze, would be beneficial.

Author response:
As we mention at the beginning of Section 4, we only expect our pre-training approach to provide benefits in the case of medium datasets (or more realistically, human datasets) where there’s a benefit to first imitating the behaviour policy, but still room for improvement. For random data, pre-training leads to an uninformed initialisation that does not provide any efficiency benefits, while for expert data, pre-training already provides optimal performance, so there is no benefit from subsequent reinforcement learning (both of which we verified empirically during our initial research).

However, following your suggestion, we have now added results from the medium-replay datasets in Appendix A.5. We find similar performance improvements across environments as for the medium-datasets, demonstrating that it is still possible to learn useful information from pre-training on mixed policy datasets.

Reviewer comment (Presentation):
Revise the structure and refine the writing to make the conclusions and emphasized information more evident.

Author response:
We have restructured and re-written Sections 4 and 5 following your advice to better separate our experimental setup and implementation details from the results and discussion in each section. Each paragraph now corresponds to a separate point of analysis, and we have highlighted the key conclusions in each section. We have also renamed Section 5 from ‘Application to Adroit Environments’ to ‘Extension to Data-Limited Adroit Environments’, and added a Section ‘Motivation for Actor and Critic Regularisation with Pre-Training’ to explain the necessary extension of our methodology to handle the limited-data setting considered in this section.

Reviewer comment (Presentation):
The two parts of the methods seem to be added incrementally rather than integrated as natural components. The authors need to improve the presentation to address this problem.

Author response:
As mentioned above, we have now added an additional Section (5.1) to better justify our addition of combined actor and critic regularisation with pre-training. To summarise here, pre-training trains the actor and critic networks to first match the behaviour policy distributions, which improves the efficiency of subsequent off-policy RL, as we show in Section 4. However, in data-limited domains, subsequent off-policy RL can lead to policy collapse, where the performance drops to close to zero after pre-training. To prevent this, and to try to maintain at least the pre-training performance, we can add additional regularisation to both the actor and the critic to prevent chain extrapolation away from the initialised behaviour distributions. Therefore, our proposal is to pre-train both the actor and critic to match the behaviour policy, and then regularise them to prevent extrapolation significantly out-of-distribution of the behaviour distributions.

While we provide this additional justification of additional regularisation in Section 5.1, we also anticipate and highlight this idea throughout the paper, with the addition of regularisation being mentioned in the last paragraph of the introduction (Section 1), the overview of our approach in Section 3 (including Algorithms 1 and 2), and the abstract and conclusion.

Reviewer comment (Evaluation):
The reviewer finds the significance of the MC pretraining unconvincing based on the results in Appendix A.3 and Section 5.

Author response:
Appendix A.3 provides an ablation which isolates the significance of the MC pre-training, so we are confused about why the reviewer is unconvinced of the significance based on these results. 

In Figure 2, we see that pre-training both the actor with imitation learning and the critic with a Monte-Carlo value error (blue) provides an efficiency benefit over randomly initialising these networks (black, still using additional LayerNorm). However, one might wonder if pre-training the actor is enough, since this is what leads to non-zero performance during pre-training (and does not depend on the critic values at all, since it is only trying to imitate and not optimise behaviour during this stage). 

Therefore, in Appendix A.3, we perform actor only pre-training with imitation learning. We see that during pre-training performance is equivalent as expected, but after pre-training performance drops off and takes a long time to recover, generally performing similarly to training from scratch with TD-learning as in Figure 2. Therefore, since the blue line reaches maximal performance more quickly than both the black line in Figure 2 and the red line in Figure 4 (Appendix A.3) we have isolated that the significant efficiency gain is due to our proposed MC pre-training.

In Section 5, we agree that pre-training does not help in these environments (as we mention) since performance collapses with or without pre-training. Therefore it is the combined regularisation approach that we propose that is responsible for the performance gains over the baselines in these data-limited environments. However, the pre-training stage is still beneficial for obtaining a non-zero initial baseline performance that allows for easier tuning of the regularisation parameters to prevent subsequent performance collapse.

Reviewer comment (Experiments):
How is the performance of the pretraining approach on "medium-replay" and "random" datasets?

Author response:
As we mention at the beginning of Section 4, we only expect our pre-training approach to provide benefits in the case of medium datasets (or more realistically, human datasets) where there’s a benefit to first imitating the behaviour policy, but still room for improvement. For random data, pre-training leads to an uninformed initialisation that does not provide any efficiency benefits, while for expert data, pre-training already provides optimal performance, so there is no benefit from subsequent reinforcement learning (both of which we verified empirically during our initial research).

However, following your suggestion, we have now added results from the medium-replay datasets in Appendix A.5. We find similar performance improvements across environments as for the medium-datasets, demonstrating that it is still possible to learn useful information from pre-training on mixed policy datasets.

Reviewer comment (Experiments):
It is necessary to conduct more ablation experiments with additional environments.

Author response:
We appreciate that performing all ablations on all datasets and environments would be ideal, but we do not believe this is necessary given that we have already run our most important ablations across all considered environments. 

For example, while Figure 3 in Appendix A.2 only considers one environment to separate the effect of LayerNorm on the actor and the critic, we ablate the effect of LayerNorm vs no LayerNorm on all environments in Figure 2. We break this down further in the HalfCheetah environment in Appendix A.2 because this environment was the only one which showed evidence of LayerNorm limiting convergence performance. Similarly, we only show intermediate values of lambda for one environment in Figure 5, but we show $\lambda=0/0.1$ and $\lambda=1$ (after pre-training) for all environments in Figure 4 in Appendix A.3. 

Finally, since we now have complete results for the medium-replay datasets in Appendix A.5 following your feedback, we have removed the final ablation you mention (Figure 6) which previously highlighted the empirical promise of applying our approach to mixed-policy datasets.

Reviewer comment (Experiments):
How does pretraining and the regularizations perform on the AntMaze domain?

Author response:
Following your feedback, we have now added the AntMaze (umaze) domain as an additional experiment to Section 5, since this indeed provides another suitable environment/dataset where imitation provides a non-trivial performance but this can be improved upon with off-policy reinforcement learning. We use an identical training procedure to that used for the Adroit environments, and show similar performance benefits of our pre-trained and regularised approaches, with relevant hyperparameters added to Table 3 in Appendix A.5.

Reviewer comment (Presentation):
The motivation example is rather unnecessary due to its simplicity.

Author response:
To address your concerns:

> The motivation example is rather unnecessary due to its simplicity

We believe that the motivational example is helpful to isolate the core idea and general applicability of our method before the implementation complexity of offline reinforcement learning algorithms and practicality of real environments and datasets are introduced in later sections. Reviewer NdGK also listed the motivational example as a strength of the paper to provide intuition for the method.

Reviewer comment (Theory):
The consistence of the actor and critic networks play critical roles in the proposed method, while it is unclear how much degree of consistence is needed in order to make it work well for off-policy training. If it is possible to derive any explicit criteria on this matter?

Author response:
To address your concerns:

> The consistence of the actor and critic networks play critical roles in the proposed method, while it is unclear how much degree of consistence is needed in order to make it work well for off-policy training? If it is possible to derive any explicit criteria on this matter?

While it is difficult to derive equations that apply in practice for deep reinforcement learning algorithms, our empirical ablations attempt to address this matter. Figure 4 in Appendix A.3 demonstrates that when both the actor and the critic are initialised to the behaviour policy during pre-training as we propose, off-policy training follows the blue line. However, if the actor matches the behaviour policy (pre-trained with imitation learning), but the critic is randomly initialised (not at all consistent), then off-policy training follows the red line. Therefore these lines serve as empirical ‘bounds’ representing no consistency (random critic, red) and approximate consistency (pre-trained critic, blue), which other degrees of partial consistency should lie between.

Thank you again for your helpful review. We hope we have addressed your questions and you can appreciate the value of our contribution towards understanding and improving the efficiency of offline reinforcement learning.

Reviewer comment (Evaluation):
What is the main reason for the huge performance drop in Walker2d-edium-EDAC?

Author response:
> What is the main reason for the huge performance drop in Walker2d-edium-EDAC?

Fundamentally, this arises because at the end of pre-training we change the objective of the actor from trying to choose the action that would have been chosen in the dataset (imitation learning), to trying to choose the action that will maximise the return i.e. critic prediction (off-policy reinforcement learning). If the values predicted by the critic are sufficiently accurate for the behaviour policy as a result of our proposed critic pre-training, then the performance should smoothly improve as we see in the Hopper environment. However, if the values are incorrect then the performance will drop, as we see in the extreme case for the red lines in Figure 4 in Appendix A.3 where the critic is not pre-trained at all.

In the Walker2d environment, we have an intermediate case where the values have been pre-trained, but are high enough variance that sometimes an action is selected which is erroneously expected to have a high value but it does not, leading to a performance drop before the critic is updated to correct its value prediction. This is an example of the bias-variance tradeoff, as we discuss in Section 3.2. The values are generally higher variance for EDAC than for TD3+BC because the critic must predict soft Q values, which also incorporate an entropy bonus as defined by Equation 5. We believe the HalfCheetah and Walker environments are higher variance because the medium trajectories end with timeouts rather than termination (as for all but one trajectories in Hopper-medium) which are more difficult to  predict without access to the current timestep [1], as we mention in Section 4. The Walker2d-medium-EDAC case therefore has the highest variance because it uses soft Q-values incorporating entropy, the episodes end with a difficult to predict timeout condition, but the episodes can still end in termination (leading to high variance on evaluation). 

To mitigate this performance drop, we can add in a temporal-difference component as in Equation 3, which reduces the variance at the cost of bias - in our case reducing the performance drop at the cost of efficiency. While we use a small temporal difference component for these environments ($\lambda=0.1$), for increasingly large datasets used for offline RL, we generally expect that the variance will be lower and a temporal difference component will not be required for pre-training.

References:

[1] Time Limits in Reinforcement Learning, https://arxiv.org/abs/1712.00378

Reviewer comment (Evaluation):
What is the computation complexity for the pretraining phase? Does it exceed a lot comparing with the offline training?

Author response:
> What is the computation complexity for the pretraining phase? Does it exceed a lot comparing with the offline training?

The pretraining phase has a lower computational complexity than the offline RL training (for the case that $\lambda=0$). This is because each pre-training update requires one forward and backward pass of the actor and critic networks to regress them towards their static supervised target. However, each offline RL training update requires an additional forward pass of the actor network for the subsequent timestep and then the target critic on the subsequent timestep in order to compute the temporal difference target ($Q(s_{t+1}, \pi(s_{t+1})$) for the critic to regress towards. This means each pre-training update is also generally faster than each offline RL update. If $\lambda \neq 0$, then the computational complexity is effectively equivalent.

Reviewer comment (Evaluation):
The pretraining stage increases the complexity of the overall training pipeline. From the training curve, we can see that the training converges faster than that without the pertaining, but almost for each environment, the performance would first drop and then begin to improve, which is quite weird in terms of robustness for the training.

Author response:
This is indeed an interesting phenomenon. Fundamentally, this arises because at the end of pre-training we change the objective of the actor from trying to choose the action that would have been chosen in the dataset (imitation learning), to trying to choose the action that will maximise the return i.e. critic prediction (off-policy reinforcement learning). If the values predicted by the critic are sufficiently accurate for the behaviour policy as a result of our proposed critic pre-training, then we expect the performance to smoothly improve as we see in the Hopper environment. However, if the values are inaccurate then the performance may drop, as we see in the extreme case for the red lines in Figure 4 in Appendix A.3 where the critic is not pre-trained at all.

In the HalfCheetah and Walker environments, we have an intermediate case where the values have been pre-trained, but are high enough variance that sometimes an action is selected which is erroneously expected to have a high value but it does not, leading to a performance drop before the critic is updated to correct its value prediction. We believe the HalfCheetah and Walker (medium) environments are higher variance because the medium trajectories end with timeouts rather than termination (for all but one trajectory in Hopper-medium) which is more difficult to predict without access to the current timestep [1], as we mention in Section 4. This is an example of the bias-variance tradeoff, as we discuss in Section 3.2. To mitigate this performance drop, we can add in a temporal-difference component as in Equation 3, which reduces the variance at the cost of bias - in our case reducing the performance drop at the cost of efficiency. While we use a small temporal difference component for these environments ($\lambda=0.1$), for increasingly large datasets used for offline RL, we generally expect that the variance will be lower and a temporal difference component will not be required for pre-training.

Reviewer comment (Evaluation):
Please explain the performance drop phenomenon in detail.

Author response:
This is indeed an interesting phenomenon. Fundamentally, this arises because at the end of pre-training we change the objective of the actor from trying to choose the action that would have been chosen in the dataset (imitation learning), to trying to choose the action that will maximise the return i.e. critic prediction (off-policy reinforcement learning). If the values predicted by the critic are sufficiently accurate for the behaviour policy as a result of our proposed critic pre-training, then we expect the performance to smoothly improve as we see in the Hopper environment. However, if the values are inaccurate then the performance may drop, as we see in the extreme case for the red lines in Figure 4 in Appendix A.3 where the critic is not pre-trained at all.

In the HalfCheetah and Walker environments, we have an intermediate case where the values have been pre-trained, but are high enough variance that sometimes an action is selected which is erroneously expected to have a high value but it does not, leading to a performance drop before the critic is updated to correct its value prediction. We believe the HalfCheetah and Walker (medium) environments are higher variance because the medium trajectories end with timeouts rather than termination (for all but one trajectory in Hopper-medium) which is more difficult to predict without access to the current timestep [1], as we mention in Section 4. This is an example of the bias-variance tradeoff, as we discuss in Section 3.2. To mitigate this performance drop, we can add in a temporal-difference component as in Equation 3, which reduces the variance at the cost of bias - in our case reducing the performance drop at the cost of efficiency. While we use a small temporal difference component for these environments ($\lambda=0.1$), for increasingly large datasets used for offline RL, we generally expect that the variance will be lower and a temporal difference component will not be required for pre-training.

Reviewer comment (Evaluation):
Please provide the training curves on hard environments such as Adroit.

Author response:
We have now provided these training curves in Appendix A.7. There is no noticeable drop-off in performance in these environments. However, as we discuss at the beginning of Section 5, we initially found that the pre-training performance collapsed entirely on these environments due to the fact that the datasets are very limited, so the critic values were leading to actions that were very different to those found in the dataset and therefore collapsing performance as we discuss in Section 2.2. As a result of this initial performance collapse, we added additional regularisation to ensure both the actor and critic remained close to the behaviour policy and prevent the actor capitalising on the extrapolated critic variance. This prevents the actor from taking actions that are substantially different to the behaviour policy and therefore mitigates any performance drop after pre-training.

Reviewer comment (Experiments):
Please explain why the pertaining steps are different for different environments, is that a hyperparameter?

Author response:
Yes, the number of steps of pre-trainining can be considered a hyperparameter. However, since both the actor and critic objectives in this pre-training phase are supervised objectives fit with a mean squared error, the loss is smooth and convex, so it is straightforward to see when these losses have converged to move onto the subsequent offline reinforcement learning objective.

Reviewer comment (Experiments):
I wonder if the pertaining loss can be a regularization loss in additional to previous regularization loss, maybe in this way we can make sure the whole training is more robust and the training curve would be more smooth.

Author response:
Yes, we agree that this could be sensible and we tried using the Monte-Carlo values as regularisation ourselves at one point during our research. However, we found that the performance can still drop using these values, and the standard CQL loss regularisation was more effective at preventing performance drop due to large OOD Q values, so we applied this to TD3+BC to regularise the critic instead. This also makes the training curves more smooth, as now provided in Appendix A.7.