Paper: On Trajectory Augmentations for Off-Policy Evaluation (ICLR-2024)

Reviewer comment (Presentation):
The main paper omitted some significant parts and placed them in the appendix, however should be explained in the main context, such as how the latent latent prior is constructed. And how the C different clusters are initially divided when identifying the PST is not clear.

Author response:
Due to the space limitations, we leave the technical details straightforwardly adapting VAE into MDP in Appendix. But **we moved the details of constructing latent prior to Section 2.2 (highlighted in blue)**. The number of clusters is determined by a fully data-driven procedure following the original TICC paper (i.e., C is determined with the highest silhouette score in clustering historical trajectories) [1].

Reviewer comment (Evaluation):
Even though the experiment could show that the OPE method performance is improved, the paper is suggested to make a fair analysis of how reliable the reward is from the algorithm augmentation, which, however, is significant for the result value of off-policy methods.

Author response:
We provided Tables 4-16 in Appendix C.1, presenting the original results, in terms of four different metrics (i.e., mean absolute errors (MAE), rank correlation, regret@1, and regret@5), on each OPE method without and with augmentation in Adroit. Given MAE directly measures the distance between estimated and true returns, we can observe that OAT leads to more accurate return estimates for different OPE methods on target policies. Especially for IS that purely relies on the importance ratio over behavior and target actions and rewards from historical trajectories, OAT effectively facilitates it on different tasks, indicating OAT providing more reliable rewards. Moreover, **we added scatter plots plotting the true returns of each target policy against the estimated returns, in Appendix A.2.1.** We can observe that the estimated rewards of each OPE method were generally improved by OAT, with less distance to true rewards.

Reviewer comment (Writing):
As shown in Fig 3, it is suggested to introduce how the latent prior is constructed, since it is the key step in constructing the new augmented PST.

Author response:
Thanks for the suggestion. **We moved the details of constructing latent prior to Section 2.2 (highlighted in blue).**

Reviewer comment (Experiments):
What is the complexity of the training VAE-MDP? Is there any analysis conducted to show the relations between training time and TDSS length/Potential trajectory scale?

Author response:
Given that training VAE-MDP requires stochastic gradient descent algorithms with parameter tuning (e.g., for step size and network architecture), we found that analyzing the complexity theoretically would be challenging, as there lacks a standard framework to facilitate that. Instead we provide additional empirical results here on training time -- specifically, **we plotted relations between training time and length of PSTs/corresponding TDSSs, in Appendix A.1.1.** The results show that the training time of VAE-MDP is increased almost linearly with the length of PSTs.

Reviewer comment (Evaluation):
In Fig 4, it is obvious that the trajectories coverages are different on the left hand and right hand corner on the top of the maze, however if these two corners are not the potential area selected to augment, should they maintain the original distribution to the maximum extent? Apart from this illustration, it is suggested that the paper provides quantitative evaluation of the difference between augmented trajectory and ground truth trajectory.

Author response:
Figure 4 is used as a toy example to illustrate the motivation behind PSTs mining and augmentation. To be more precise, the augmented sub-trajectories are not necessarily temporally-aligned, i.e., the starting time and ending time of found PSTs on each trajectory can be varied as described in Appendix C.3. As a result, a small number of sub-trajectories may cover some parts of the left-top and right-top corners (which makes it particularly obvious when augmenting 10 times numbers of trajectories in Figure 4). Nevertheless, one can still observe that augmented PSTs are effectively improving state-action coverage. 

Since our experimental environments (e.g., Adroit) are more complex than the Maze2D-umaze environment (state space is roughly $\mathbb{R}^{20}$), it would be very challenging to quantify the distance between augmented and original trajectories. Please note that we had conducted extensive experiments over 34 datasets over multiple application domains, encompassing varied levels of state-action coverage, human involvements, horizons, density of rewards, dimensions of states and actions, etc. Moreover, we used 4 different evaluation metrics to thoroughly examine the performance of 5 OPE methods based on the augmented trajectories.  From the results, one can observe that OAT effectively enhances all types of OPE methods, as in Appendix C.1, illustrating the importance of augmenting trajectories for OPE to better estimate target policies. Moreover, we added scatter plots plotting the true returns of each policy against the estimated returns, in Appendix A.2.1. We can observe that the estimated rewards of each OPE method were generally improved by OAT.

Reviewer comment (Writing):
It feels like the paper was somewhat rushed. There are some confusing parts/writing errors such as: what is "human-involved RLs" (third line in intro), "while human may behave" (5th line last para of intro), Figure 2 I believe should be referencing Sections 2.1/2.2/2.3 (not 3.X), "intrinsic nature that human may follow" (2nd line Section 2.1).

Author response:
Q1. There are some confusing parts/writing errors such as: what is "human-involved RLs" (third line in intro), "while human may behave" (5th line last para of intro), Figure 2 I believe should be referencing Sections 2.1/2.2/2.3 (not 3.X), "intrinsic nature that human may follow" (2nd line Section 2.1).

A1. Thanks for pointing out the typos. We have fixed the references on Figure 2 in the manuscript.
For the mentioned “confusing parts’’, please see our comments point-by-point below:
* Human-involved RLs generally refer to the context of RL with human involvements, including human-involved decision-making processes, human-guided RL algorithms, etc., which are commonly discussed in related works [1-5].
* The phrase, “while human may behave”, is discussed under the context that human-involved systems have limited state-action coverage, which can be caused by homogeneous behavior policies (as discussed in 1st paragraph of Introduction); on the other hand, human can behave diversely under different policies [6-7]. Following different policies can result in varied state-action coverage [8-10], which is a fundamental problem in RL.  
* The phrase, “intrinsic nature that human may follow”, is under the context discussing an intrinsic nature/characteristic in human-involved systems, that human could follow homogeneous behavioral policies or specific guidelines when they perform their professions.

We have added more references in our manuscript, for readers' potential interests.

Reviewer comment (Writing):
The notation in Section 2.1 is very difficult to parse. I suspect there is an easier way to explain this. Also Figure 3 is very confusing, and is not explained in the text (there are some references in the appendix, but I think there should be much more explanation of it given the complexity of the diagram).

Author response:
Q2. The notation in Section 2.1 is very difficult to parse. I suspect there is an easier way to explain this. Also Figure 3 is very confusing, and is not explained in the text (there are some references in the appendix, but I think there should be much more explanation of it given the complexity of the diagram).

A2. We understood that the notations/equations may be a bit intense in section 2.1, as we are the first work that introduces to OPE with sub-trajectory augmentation, which is a relatively new framework. As a result, we chose to present our methodology as detailed as possible so readers with different backgrounds can get the idea thoroughly (in case they would like to build on top in the future). We have double checked that there do not exist any sub-/super-scripts or variables that are redundant. **We further noticed that the other reviewer pod6 specifically pointed out that this part is easy to follow, as well as WbKh who pointed out our figures are informative -- we would greatly appreciate it if the reviewer can point out the specific part of methodology that is redundant/hard to follow.**

Reviewer comment (Writing):
I think drawing connections to how humans behave (such as in Figure 2 caption) is too strong. It is appropriate for intuition purposes but making general statements on how humans behave seems a bit strong.

Author response:
Q3. I think drawing connections to how humans behave (such as in Figure 2 caption) is too strong. It is appropriate for intuition purposes but making general statements on how humans behave seems a bit strong.

A3. When the trajectories are collected from humans, the states and/or actions are highly  related to human behaviors. Similar sub-trajectories may exhibit similar human behaviors. Such findings are reported in related works [21-23].

Reviewer comment (Theory):
Biggest concern. I dont think the paper gives a good intuition for why this works. While coverage is important, it seems like that is insufficient. Another challenge is distribution shift. Prior work [1] has discussed that under extreme distribution shift, OPE methods can diverge. However, this work does not discuss this. Moreover, it seems possible to me that this data augmentation technique could introduce samples that worsen the distribution shift, which could worsen accuracy of OPE methods.

Author response:
Q4. I don't think the paper gives a good intuition for why this works. 
(Q4-1) While coverage is important, it seems like that is insufficient. 

A4-1. We respectfully disagree with the comment. Many existing works have justified that improving state-action coverage is an important and still open-ended problem [4, 8-10]. **Our work is the first to carry it onto the OPE domain, and attempts to solve it through offline trajectory augmentation.**

(Q4-2) Another challenge is distribution shift. Prior work [1] has discussed that under extreme distribution shift, OPE methods can diverge. However, this work does not discuss this. 

A4-2. We agree distribution shift is a challenge for OPE. Although the mentioned work by Wang et al. only considers a very specific type of OPE method under a specific setting (i.e., FQE with linear approximation using pre-trained features from neural networks), DOPE benchmark [10] also found that OPE methods in general have such issues. *This is beyond the scope of this work, as we use existing OPE methods as backbones to process the augmented trajectories.* Hope in the future, the communities will come up with new OPE methods that can resolve the distribution shift in OPE.

(Q4-3) Moreover, it seems possible to me that this data augmentation technique could introduce samples that worsen the distribution shift, which could worsen accuracy of OPE methods.

A4-3. As discussed in Limitations, latent-model-based models such as VAE have been commonly used for augmentation in offline RL, while they generally rarely come with theoretical error bounds provided, such as Dreamer [12, 13], PlaNet [14], SLAC [15], LatCo [9], Solar [16], etc. Such a challenge also remains for many data augmentation methods [17]. However, once the trajectories are augmented, one can choose to use the downstream OPE methods which come with guarantees, such as DR and DICE. Moreover, please note that our experiments include two real-world environments (e-learning and healthcare) where our method significantly outperforms existing works. And we conducted extensive experiments over 34 datasets with varied characteristics including application domains, state-action coverage, human involvements, horizons, density of rewards, dimensions of states and actions, etc. The experimental results show our method can generally enhance different OPE methods.

Reviewer comment (Evaluation):
Could one generate better samples by somehow tying in the evaluation policy behavior into the data generation process? It seems like if the goal is to evaluate some policy $\pi_e$, we should account for what $\pi_e$ does/samples actions?

Author response:
Q5. Could one generate better samples by somehow tying in the evaluation policy behavior into the data generation process? It seems like if the goal is to evaluate some policy \pi_e, we should account for what \pi_e does/samples actions?

A5. In this work, we chose to use behavior data only for augmentation, so that the proposed method can be **stand-alone** to generate trajectories without any assumptions over target policies. And it can be easily utilized by broader built-on-top works such as policy optimization and representation learning. We agree that augmenting target policies is a possible future direction.

Reviewer comment (Reproducibility):
How is step 1 in the algorithm actually done? That is, under what basis are states grouped into clusters. Some representation learning work does this based on behavior similarity [1] etc. How the states are grouped is important for all the remaining steps.

Author response:
Q6. How is step 1 in the algorithm actually done? That is, under what basis are states grouped into clusters. Some representation learning work does this based on behavior similarity [1] etc. How the states are grouped is important for all the remaining steps.

A6. We appreciate the reviewer pointing us to the representation learning paper. As provided in Appendix A.3, we leverage TICC problem [18] to cluster states, which captures graphical connectivity structure of both temporal and cross-attributes information and has shown . Specifically, in experiments, we utilize its variant, MT-TICC [19], which further considers time-awareness and multi-trajectories based on TICC. **We also provided detailed formulations of the TICC problem in Appendix A.3.1 (highlighted in blue).**

Reviewer comment (Theory):
The "grounding" process of Eqn 2 is unclear to me. How do you avoid the model from hallucinating and generating samples that cannot occur in the MDP since interaction with the MDP does not happen in the OPE setting?

Author response:
Q7. The "grounding" process of Eqn 2 is unclear to me. How do you avoid the model from hallucinating and generating samples that cannot occur in the MDP since interaction with the MDP does not happen in the OPE setting?

A7. This is a good question that points out a potential future avenue following our work. We took a similar approach as existing works who also adapted VAEs to learn environmental dynamics [9, 12-16, 20], which cannot resolve such a limitation, but it is definitely worth looking into.

Reviewer comment (Evaluation):
The unsupported justifications of the results — In more than one occasion, the authors (rightly) discuss very specific reasons that could confirm/justify the observed results. However, in most cases in this work, the justifications are not supported by data. For example, the authors say: "In contrast, the historical trajectories induced from simulations tend to result in better coverage over the state-action space in general, and the augmentation methods that do not consider the Markovian setting may generate trajectories that could be less meaningful to the OPE methods, making them less effective." It is unclear what is meant by "less meaningful" and whether the authors believe that non-Markovian augmentation is worse in scenarios with better state coverage more generally. The authors often use phrases such as "we conjecture" and "a possible reason" without strictly backing up the claims with evidence, which would greatly improve the quality of the paper.

Author response:
A1. **We added related citations and detailed evidence to justify our findings and statements in the revised manuscript (highlighted in blue).**

Reviewer comment (Experiments):
The lack of ablations to validate the proposed innovations — While the authors compare against a wide set of benchmarks, it is unclear whether these benchmarks are explicitly tailored for the OPE problem or not. Furthermore, since the proposed framework is composed of multiple smaller contributions (such as discretization, VAE architecture, etc.), the authors should isolate each of these contributions individually and support their claims with evidence and experiments.

Author response:
A2-1. To the best of our knowledge, no prior work systematically investigates trajectory augmentation for facilitating OPE, our work is the first to investigate it and provide a possible solution (i.e., OAT). Thus, we examine a broad range of possible data augmentation methods, which were original proposed towards *data that may sharing some characteristics to offline trajectories*, or *sharing some intuitions with OAT*, including: 4 RL-oriented (TDA, Permutation, jittering, scaling), 2 generative models (TimeGAN, VAE), 2 time series-oriented methods by capturing similar behaviors (SPAWNER, DGW). Moreover, we investigated 5 popular OPE methods that are broadly used in prior works [1-11], spanning 3 major categories of OPE as defined by Yue’s group [11].  

A2-2. Given the key idea is to mine potential sub-trajectories and considering Markovian nature on trajectories, we compared OAT to several ablations:
* OAT w.o. Discretization & support determination (TDA)
* OAT w.o. PSTs mining (VAE-MDP)
* OAT w.o. MDP (VAE)

Note that both discretization and support determination are two common successive techniques to extract shared high-level representations (e.g., PSTs) from high-dimensional complex data [14-16]. Thus, we ablated the proposed PSTs mining from two perspectives: ablating the support determination to identify PSTs from TDSSs by randomly selecting sub-trajectories to augment (which aligns with the original idea of TDA [17]); ablating the concept of PSTs by applying VAE-MDP on entire trajectories (VAE-MDP). The third ablation is applying VAE on entire trajectories, without adaptation to Markovian setting.

Reviewer comment (Experiments):
The relatively narrow-scoped experiments — The experiments have a limited scope, and it is unclear whether the authors considered applying their approach within a reinforcement learning (RL) context. There is a question about how well the proposed method would perform in such a context.

Author response:
A3. In this work, we specifically focus on the OPE problem [1-10] as the improvements resulting from the model we designed can be isolated, as opposed to policy optimization where the policy improvement step will also impact the performance. We followed the guidelines and standardized procedures introduced in a recent benchmark, DOPE [1], from Levine's group, which provided for each D4RL environment the target policies to be evaluated as well as an off-policy training dataset. The effectiveness and robustness of OAT was extensively validated over **34 datasets with varied characteristics** including applications, state-action coverage, human involvements, horizons, density of rewards, dimensions of states and actions, etc. The experiments contained 2 real-world applications, education and healthcare. Moreover, OAT can be **stand-alone** to generate trajectories without any assumptions over target policies. And it can be easily utilized by built-on-top works such as policy optimization and representation learning. **We added such discussions in the conclusion section -- highlighted in blue.**