Paper: Simple Data Sharing for Multi-Tasked Goal-Oriented Problems (ICLR-2024)

Reviewer comment (Experiments):
As the authors themselves outline in the conclusion, only AntMaze was considered. While I appreciate that the authors were open about the limited experiments, I think comprehensive experiments are something necessary to establish the strength of SDS rather than it being a “interesting future direction.” As a comparison, MAHALO considers more settings on MuJoCo (D4RL datasets) as well as MetaWorld.

Author response:
We chose AntMaze because it is a popular benchmark and it is easy to adapt its environments to construct offline CGO problems. (On the other hand, locomotion mujoco problems in D4RL are not goal oriented). . Besides the 3 variations, we also add goal prediction as an extra baseline and experiment with new perturbation settings (Section 4.4 and 4.5) when constructing the context-goal dataset, which helps to show the efficacy of the method in a better way.

Reviewer comment (Experiments):
The introduction motivates SDS by claiming that Offline GCRL algorithms such as GoFar (Ma et al., 2022) “can fail when the predicted goal is not reachable from the initial state.” However, this claim is also not evaluated in the experiments. Overall, there are a lack of baselines which were included in the related work section but not compared in the experiments.

Author response:
Thanks for your suggestion! **We also add the goal prediction baseline in Section 4.4 and 4.5.**

Reviewer comment (Experiments):
Similarly, MAHALO was not considered as a baseline for the experiments although it was compared to SDS in the theoretical results.

Author response:
We did initial trials using MAHALO in the AntMaze environments, and we found that the performance is not as good as PDS+IQL or goal prediction + IQL in larger mazes. One reason may be that the public MAHALO implementation is based on ATAC and PSPI, which do not have published results and hyperparameters for  Antmaze. Also, MAHALO requires adversarial training to learn a pessimistic reward function which might pose extra difficulty in the training.

Reviewer comment (Novelty):
While the theoretical results in Theorem 3.1 are nice, it is not immediately clear what the novel contributions are. While I am not too familiar with the literature on statistical RL and I may be missing something, I think the authors should be explicit about what novel techniques are used to extend theorems from previous work, especially if the theoretical results are the main contribution of the paper.

Author response:
The main technical contribution here is to show that SDS which is based on a fictitious action can be implemented with PSPI by using function approximators and their assumptions based on the original MDP without the fictitious action. The details are given in Appendix B.3 and Appendix B.4 ( Lemma B.10): We show that by carefully analyzing the Bellman equation of the action-augmented MDP, we can translate the needed assumptions back those defined in the original MDP, so we have Assumption 3.2-3.4 defined in natural terms wrt the original MDP. These techniques are novel and critically leverage the structure of CGO and SDS. The statistical analyses in Appendix B.4 (High-probability Events) are standard

Having the assumptions made in terms of the original MDP (not the augmented MDP) is important to establishing a formal theoretical result. Otherwise, the assumptions may be vacuous as they cannot be compared with others who do not introduce the augmented MDP.  We will highlight the importance of these contributions better in the final draft.

Reviewer comment (Experiments):
What is the main difficulty in obtaining negative data (context-non-goal pairs) which is described as a limitation of MAHALO? What makes the unsupervised dynamics data easier to obtain?

Author response:
We discuss why unsupervised datasets and context-goal datasets are commonly available in the introduction: In robotics, task-agnostic play data (unsupervised data) can be obtained at scale since it does not require any reward labels. Positive data (context-goal pairs) would be easy to obtain as "successful samples" as in some instruction datasets, while negative data under the CGO setting would be "failure samples" which are generally harder to obtain in common instruction datasets.

Reviewer comment (Writing):
Although some discussions are involved in the introduction, I think this paper still lacks explanations about the CGO setting. I doubt the coexistence of separate unsupervised data and context-goal data might not be so difficult to combine. For example, in the navigation setting mentioned by the authors, it is easy to derive a goal position from the actual state and vice versa. With a predefined context distribution, we can label a state accomplishing a given context as the goal state with ease.

Author response:
While what the reviewer mentioned is applicable in some special cases, such a hindsight construction is not possible for offline CGO problems in general. Here are the reasons:  **1) the test context which might not even exist in the context-goal data so we cannot directly label goal states (e.g. when the context space is continuous); and 2) the exact same states in the context-goal dataset could never appear in the unsupervised dataset.** Therefore, an example of this kind is our third experimental setup in Section 4.5, where the test context is randomly sampled so it could never appear in the training set (but potentially still under coverage through function approximation). In the revision, we also add new experimental results in our draft where we perturb the context-goal dataset **(Section 4.4 and 4.5) such that the goals in the context-goal dataset never appear in the unsupervised dataset, and show that our method still works well**. Finally, as in our footnote for the oracle function, it is only used to generate context-goal dataset for the sake of problem construction; this information is not revealed to the learning algorithms and with further perturbation of the context-goal dataset it is also not possible to  “label a state accomplishing a given context as the goal state with ease”.

Reviewer comment (Theory):
In Section 3.1, the authors see the key problem that "some \( s \in G_c \) in the \( D_\text{goal} \) dataset is also observed in the dynamics dataset." However, I speculate that the process of Algorithm 1 cannot tackle this issue as the transition of a goal state \( x \) to \( x' \) where \( x' \neq x^+ \) will still remain. The problem can only be alleviated with the mentioned approach of "equally balancing the samples \( \bar{D}_\text{dyn} \) and \( \bar{D}_\text{goal} \)" (in section 3.2).

Author response:
The claim that “Algorithm 1 cannot tackle this issue as the transition of a goal state x to x’ where will x’ \neq x^+ still remains” is a misunderstanding, **since a^+ is not in the original action space**. For example, x is a goal state, x’ is some non-goal state and there is some transition (x,a,x’) in the original dynamics dataset, and we add (x, a^+, x^+) as fictitious transitions. As a result, Q(x,a^+) would be high while Q(x,a) will be low, so V^*(x) will still be high. In general, we introduce an extra fictitious action a^+ for x to reach x^+, which is outside the original action space in the unsupervised dataset, so there is no such conflict in dynamic programming when we compute the optimal value function via argmax over all possible actions (including both real and fictitious actions). If such fictitious action and transition exists with non-zero probability, it will naturally be the highest Q value over all possible actions (e.g., via expectile regression for IQL,which does not necessarily require balancing the samples).

Reviewer comment (Theory):
A confusing point is that the authors provide theoretical analysis on SDS+PSPI but do not use this variant for experiments (they use SDS+IQL instead).

Author response:
Thank you for raising this point. We choose IQL since it is a popular practical algorithm to solve AntMaze environments, making it easier to compare with the literature and other baselines. Including SDS+IQL in the experiments also shows that the idea of SDS is indeed general, not just limited to PSPI.

Reviewer comment (Experiments):
Why do the experiments replace 0/1 rewards with -1/0 rewards? Are they intrinsically different or empirically different?

Author response:
To keep our method comparable to the context agnostic baseline, we follow the tradition to use -1/0 reward in AntMaze experiments as in the original IQL paper. The -1/0 rewards are just a scaled and shifted version of the 1/0 rewards (specifically, first scale the reward of 1/0 by $\frac{1}{1-\gamma}$ to the reward of $\frac{1}{1-\gamma}$ / 0 and then shift the reward at each state by -1; then we get the -1/0 reward). Therefore these two reward systems induce the same policy ranking, and our theory applies.

Reviewer comment (Experiments):
Why not use the theoretically sound SDS+PSPI method for experiments?

Author response:
Thank you for raising this point. We choose IQL since it is a popular practical algorithm to solve AntMaze environments, making it easier to compare with the literature and other baselines. Including SDS+IQL in the experiments also shows that the idea of SDS is indeed general, not just limited to PSPI.

Reviewer comment (Reproducibility):
I'm confusing about the reward prediction approach adopted to the baselines. How can we train a reward model with only using positive data (from \( D_\text{goal} \))? For the usage of the reward model, can you elaborate the meaning of "choosing the percentile as 5% in the reward distribution evaluated by the context-goal set as the threshold to label goals" in Appendix C.1?

Author response:
As mentioned in 4.2, we use PDS to train a pessimistic reward function with ensembles of 10 reward functions, which measures the uncertainty of the value function prediction, making it possible to learn a pessimistic reward as introduced in PDS. Since the uncertainty term is subtracted from the reward predicting, making it a continuous reward prediction and we need to convert it back to -1/0 to label the goals, we need to select some certain threshold. So we use the tradition in outlier detection where we choose some percentile of the predicted values from the training distribution as the threshold. 

Nonetheless, we also note that the pessimistic rewards learned in this way, and the more naive version RP, are not perfect and can be far away from the true rewards. Nonetheless, we saw in the experiments, IQL with these imperfect rewards can still have fairly good performance in some datasets. Such a robust behavior to reward imperfection has been recently observed and proved for offline RL in the literature [Li et al. 2023, Shin et al. 2023].

Li, Anqi, et al. "Survival Instinct in Offline Reinforcement Learning." arXiv preprint arXiv:2306.03286 (2023).

Shin, Daniel, Anca D. Dragan, and Daniel S. Brown. "Benchmarks and algorithms for offline preference-based reward learning." arXiv preprint arXiv:2301.01392 (2023).

Reviewer comment (Evaluation):
In Section 4.3, what is the reason for the observation that "for PDS, we can observe that the reward distribution for positive and negative samples are better separated in the large one than the medium one?"

Author response:
For different environments, we have different context-goal dataset and trajectory only dataset. The performance of PDS relies on how well it can capture the pessimism correctly in our test context distribution given the training data. And we observe that although the medium maze is a smaller problem than the large maze, the training distribution is harder for PDS to separate the positive and negative examples for the test distribution.

Reviewer comment (Experiments):
What are main differences among the three experimental settings (from Section 4.3 to 4.5)? I observe a different range of context distribution but it seems that all test contexts are either fully in-distribution or partially in-distribution (Table 4) but not include real out-of-distribution cases.

Author response:
The main differences are shown in Fig 1 and also discussed in 4.1: 1) the easy case where goal sets are very similar such that it could be approximately solved by a context agnostic policy, 2)We have finite contexts and the goal sets from different contexts are distinct and do not overlap  3)The contexts are continuous and randomly sampled, the goal sets from different contexts are different but could share some overlap, making a complex CGO problem. We do not aim to show out-of-distribution generalization for context distributions. However, the learned policy could enjoy some out-of-distribution generalization which would be an interesting future direction.

Reviewer comment (Theory):
The theoretical assumptions seem imprecise or incorrect. For example, the authors say that the state, action, and context spaces can be continuous. If they are continuous (say $|\mathcal{A}|=|\mathcal{S}|=|\mathbb{R}|$), then $|\Pi|$ is extremely large (e.g contains $|\mathcal{A}|^{|\mathcal{S}|}=\aleph_1^{\aleph_1}$ deterministic policies). In this case $\epsilon_{dyn}$ will be extremely large, making the RHS of the bound in Theorem 3.1 extremely large.

Author response:
Thank you first for raising this point. As we discussed in the footnote of the theorem, we only use the discrete notation (i.e. writing the cardinality of function and policy classes) for simplicity and “We state a more general result for non-finite function classes in the appendix”. In Appendix B4, **we introduce the definition of covering number and replace |\Pi| with the covering number of the function class**.

Reviewer comment (Theory):
If I understand Theorem 3.1 correctly, then the right-hand side of the bound is almost always greater than $V_{max}=\frac{1}{1-\gamma}$ (since the rewards are only non-zero at goal states and are 1 otherwise). Since this would make the bound useless, I am probably misunderstanding the meaning of each term on the RHS. It would have been helpful if the authors explained step-by-step each term of the bound and how that leads to a small regret with high probability (maybe using an example domain).

Author response:
Certainly, we can interpret the upper bound in the theorem as follows: the statistical errors would decrease as we have more data from $\mu_{goal}$ and $\mu_{dyn}$; then with the finite coefficients, the final upper bound would also decrease. Take $\pi = \pi^*$ as an example. For the coefficients to be finite,  it means 1) the state-action distribution from the dynamics data ``covers" the state-action pairs generated by $\pi^*$, which naturally includes stitching; 2) the support of $\mu_{goal}$ "covers" the goals $\pi^*$ would reach. The coverage is measured based on the generalization ability of $f$ and $g$; e.g., if $g(x_1)$ and $g(x_2)$ are similar for $x_1\neq x_2$, then $x_2$ is within the coverage so long as $x_1$ can be generated by $\mu_{goal}$.

Reviewer comment (Writing):
The presentation of the theoretical results is very poor (mainly Section 3.3). There is barely any explanation of what the Theorem, Assumptions, and Definitions are saying. For example, Definition 3.4 defines "generalized concentrability" out of the blue and does not explain what it is saying and how it is relevant. It is also not clear why it uses the same symbols that were defined for just "concentrability" in Theorem 3.1.

Author response:
Thanks for your suggestions. We reorganized Section 3 accordingly (e.g., explanations, definitions). Please check the revised draft.

Reviewer comment (Experiments):
All the experiments are in a single domain (the D4RL AntMaze). While the 3 variations in goal distributions are useful, comparisons with baselines in other domains (e.g. another D4RL domain) would have helped get a better sense of the general applicability of the proposed approach.

Author response:
We chose AntMaze because it is a popular benchmark and it is easy to adapt its environments to construct offline CGO problems. (On the other hand, locomotion mujoco problems in D4RL are not goal oriented). . Besides the 3 variations, we also add goal prediction as an extra baseline and experiment with new perturbation settings (Section 4.4 and 4.5) when constructing the context-goal dataset, which helps to show the efficacy of the method in a better way.

Reviewer comment (Theory):
The proposed approach is only applicable to goal-conditioned tasks with rewards of 1 only at goal states and zero rewards otherwise.

Author response:
We agree such goal oriented problems do not cover the full spectrum of RL problems, but they are quite common in practice, especially given many real-world problems do not come with dense rewards naturally. That is why goal conditioned RL (GCRL) problems are widely studied in the literature. We choose the sparse reward (only 1 at goal states and 0 otherwise) setting as it is a realistic and commonly-used setting in goal conditional RL where the reward is sparse. It is studied in most of the GCRL papers [1]. We should note that this setting also covers problems generated by applying any reward shifting/scaling (e.g.,GCRL problem with -1 at each step and 0 at the goal) since they will induce the same policy ranking.

[1] Goal-conditioned reinforcement learning: Problems and solutions. M Liu, M Zhu, W Zhang

Reviewer comment (Theory):
The paper makes very strong assumptions about the offline dataset. Mainly, they assume that the goals dataset ($D_{goal}$) contains all the goals the agent will encounter and that the trajectories dataset ($D_{dyn}$) contains trajectories leading to those goals.

Author response:
The assumption stated here is indeed a misinterpretation of our theory. We are sorry if our original manuscript’s writing is not accurate. To clarify, we only require 1) the distribution generating the dynamics dataset "covers" a feasible path to the target goal set and 2) the context-goal distribution "covers" some feasible goals. We remark that "covering'' here does not mean that such a path or a goal has to appear in the datasets or the data have to be generated only by goal-reaching policies (e.g., it suffices that data can be stitched together to form a goal-reaching path). Instead, it is a distribution-level coverage requirement; especially **we present coverage conditions based on the generalization ability of function approximators used in learning, which is weaker than requiring coverage conditions based on non-zero density**. We also add additional explanations in our assumptions and theorems in the draft.

Reviewer comment (Experiments):
My concern about the experiments remains. I think the authors should also demonstrate their approach in other domains.

Author response:
We chose AntMaze because it is a popular benchmark and it is easy to adapt its environments to construct offline CGO problems. (On the other hand, locomotion mujoco problems in D4RL are not goal oriented). . Besides the 3 variations, we also add goal prediction as an extra baseline and experiment with new perturbation settings (Section 4.4 and 4.5) when constructing the context-goal dataset, which helps to show the efficacy of the method in a better way.

Reviewer comment (Theory):
Theory 3 does not address my concern that the reward scheme used in Algorithm 1 is potentially problematic. If $(s_0,a_0,s_1),...,(s_t,a_t,s_{t+1}),...,(s_{T-1},a_{T-1},s_T)$ is a trajectory in $D_{dyn}$ and $(s_t,c)$ is in $D_{goal}$, then $\bar D_{dyn}$ will contain $((s_t,c), a^+, 0, (s_{t+1},c))$ but $\bar D_{goal}$ will contain $((s_t,c), a^+, 1, (s^+,c))$.

Author response:
We present why our augmented MDP is equivalent to the original MDP in the sense of regret in Appendix B1, and add Remark 3.1 as an intuition. The core idea is that by introducing an extra “fictitious” action $a^+$, we create some fictitious transition from the potential goal states to the terminal state as future transitions after reaching the goal state. **For such potential goal states, the positive future reward will be propagated back to them with a high value by Bellman equation.** For example, if s0->...->g is a trajectory to some potential goal g, adding future transition g - > s+ with action a^+ and reward one, and labeling all previous trajectories with reward 0 will not affect the ranking of the learned value functions.