Offline Model-Based Skill Stitching

Weaknesses:

1. The assumption of the availability of a dataset for each skill is a strong assumption, is there a way to relax it? For example learning diverse skills from one offline dataset? Is this possible and is there any related work that focus on this problem?

We are not sure what the "one dataset" here exactly refers to.

• For one offline dataset with various single skills, we can divide it into parts of each skill.
• For one offline dataset with only long-horizon tasks containing many skills in order, we can divide it into parts of each skill.
• However, for one offline dataset with only one single skill, it might be unsolvable, since we have no access to the trajectories of other skills. The policies, value functions, and dynamics models trained on datasets of only one skill, are possibly unable to generalize to out-of-distribution states and transitions.

2. Training each skill separately via offline RL seems expensive and time-consuming.

We train a multi-task policy instead of multiple sing-task policies in both Maze and Kitchen, which can be viewed as a goal-conditioned policy. Take Kitchen as an example. Following the tradition of vector-based RL, we use a 7-dimensional one-hot vector (since there are 7 tasks in Kitchen) to denote which skill is being executed, and the one-hot vector is concatenated with the observation vector before taken by the policy network as input. We have made this clear in the revision (modified texts in blue in Section 4.1.2).

3. For some hyperparameters it is not clear to me they have been chosen, for example the maximum steps of skill execution seems very problem dependent.

The maximum steps for each environment are independent. To be concrete, in the Maze environment, the maximum steps for a task containing 2 / 3 / 4 skills are all 50, including 5 steps in maximum of each stitching part between two adjacent skills. In the Kitchen environment, the maximum steps for 2 / 3 / 4 skills are 250 / 360 / 480 respectively, including 20 steps in maximum of each stitching part between two adjacent skills. In the Robosuite environment, the maximum number of step for open door $\to$ close door is 50, and the maximum number of step for can $\to$ can is 100, including 20 steps in maximum of each stitching part between two adjacent skills. Note that we utilize action spaces referred from MAPLE, which greatly reduces the number of steps needed in the Robosuite benchmark compared to previous literature. We have made this clear in the revision (modified texts in blue in Appendix B.3).

4. The method does not seem effective on more complicated tasks (for example in table 2 the method fails in accomplishing more than one skill regardless of the number of skills in the task), but it is still better than the baselines.

This is highly due to the low success rate of the first skill, which is not caused by our stitching method. For example, the success rates of the 7 tasks in Kitchen are listed below:

Since our experimental results are based on the average over all possible permutations, the absolute value of scores will be inevitably low. For example, consider the first row of Table 2 in the paper, where we test over $A_7^2=42$ combinations of tasks. Among the final scores of these 42 tasks, at least $4\times 6=24$ of them will be 0, since 4 of the 7 tasks cannot be accomplished at all. Theoretically, the upper bound of the score accomplished by an arbitrary (even expert) algorithm should be $0.96 \times 2 \times \frac{1}{7} + 0.00 \times 2 \times \frac{4}{7} + 0.88 \times 2 \times \frac{1}{7} + 0.27 \times 2 \times \frac{1}{7} = 0.603$ (without normalize, as the revised version has normalized the results according to another reviewer). As shown in Table 2 (before revision), Our MB-Stitch has achieved 0.518, which is close to the upper bound, and outperforms the baselines.

Weaknesses:

1. Lack of novelty:

[1] trains a critic $Q(s, a)$ with goal-relabeled short sub-trajectories, then trains a value-conditioned diffusion model to generate trajectories given a goal, and finally executes the actions in the generated trajectories. [2] introduces trajectory stitching for synthesizing new trajectories (as a method of data augmentation), and trains BC policies over augmented datasets. Although [3] uses model-based rollouts (planning) for skill-based task planning, it has less to do with skill stitching.

Unlike previous literature [1, 2], our paper lies in a totally different problem setting, where we regard each task as a skill. [1] and [2] both stitch sub-trajectories within a single task, while we consider task-level stitching in our paper. When generating an extra sub-trajectory $\tau_2$ to be stitched after the former sub-trajectory $\tau_1$, there is no gap between $\tau_1$ and $\tau_2$. However, when stitching two task-level skills (e.g., microwave and kettle), the gap between two adjacent tasks is rather large, which is the biggest difference in between. In a word, our method tackles the challenge of skill stitching in the case where there are large gaps between two task-level skills, which is different from previous works [1, 2].

[1] Stitching Sub-trajectories with Conditional Diffusion Model for Goal- Conditioned Offline RL (AAAI 2024)

[2] Model-based Trajectory Stitching for Improved Offline Reinforcement Learning (Offline RL Workshop @ NeurIPS 2022)

[3] Offline Policy Learning via Skill-step Abstraction for Long-horizon Goal-Conditioned Tasks (IJCAI 2024)

Questions:

1. I wonder if the value function properly evaluates states that have not been visited (during stitching). As the value functions for each skill are learned distinctly, how can the value evaluation in the stitched space be accurate and reliable?

Learning the value functions distinctly for each skill does not mean the inaccuracy. As long as the used value function represents the relative suitability (possibility) of accomplishing the next skill (task), it can be utilized to guide the process of stitching. Besides, we adopt the trick of model ensemble to reduce the inaccuracy of value functions, which can make value estimation more reliable.

2. How might the proposed method be adapted to handle low-coverage offline datasets?

In the phase of learning individual skills, our method faces the same challenges as offline RL methods when data coverage is limited. However, in the skill stitching phase, our approach can mitigate the scarcity of each individual skill dataset by training on the union of all datasets, where the necessary transitions between skills are more likely to be present. Additionally, our model-based approach leverages the generalization capabilities of world models to fill in missing transitions, while incorporating a conservative strategy to mitigate distribution shift.

3. I wonder if the authors considered any techniques to reduce the computational burden of MPC in continuous or stochastic environment?

To accelerate MPC, techniques such as approximate planning (e.g., TD-MPC [4], which uses learned value functions to bootstrap planning), batch computation, model compression and distillation, and receding horizon control [5] could be applied. While these methods are promising, we have not focused on efficiency in this work, as our primary contribution is improving the stitching performance of offline-learned skills.

[4] Hansen et al., "Temporal difference learning for model predictive control." (2022).

[5] Mayne and Michalska. "Receding horizon control of nonlinear systems." (1988)

4. What potential strategies could be considered for improving generalization or adaptability to dynamic environments within the constraints of offline learning?

Data augmentation could be properly considered, utilizing generative models. Through data augmentation, we could get larger quantities of data with a higher coverage of states and transitions.

5. Minor typos.

Thank you for pointing out the typos! We have corrected these in red color in the revision.

Weaknesses:

1. About originality and baselines:

PEX [1], OPAL [2] and LPD [3] all tackle challenges in single long-horizon tasks. For example, PEX [1] introduces offline-trained policy $\pi_\beta$ and online-trained policy $\pi_\theta$, but these two policies are used as a policy pool to execute within a single task. During the execution phase, an action is produced from the joint policy $\Pi=[\pi_\beta, \pi_\theta]$ for each step. OPAL [2] introduces hierarchical structures to stitch primitive policies.

Different from previous literature, our paper lies in a totally different problem setting, where we regard each task as a skill. In other words, we consider task-level stitching. The largest gap between primitive-level stitching and task-level stitching is that task-level stitching requires inserting an unseen trajectory sequence, which is not needed in primitive-level stitching. Considering an example of task-level stitching: microwave $\to$ kettle, there is a giant gap between the last state of finishing the task “microwave” and the initial state to accomplish the task “kettle”. However, for a primitive-level case such as “kettle = move + grasp + lift”, there is no gap when switching from “grasp” to “lift”. In a word, our method tackles the challenge of skill stitching in the case where there are large gaps between two task-level skills, which is different from previous works [1, 2, 3].

[1] Zhang, Haichao, We Xu, and Haonan Yu. "Policy expansion for bridging offline-to-online reinforcement learning." arXiv preprint arXiv:2302.00935 (2023).

[2] Ajay, A., Kumar, A., Agrawal, P., Levine, S., & Nachum, O. (2020). Opal: Offline primitive discovery for accelerating offline reinforcement learning. arXiv preprint arXiv:2010.13611.

[3] Yang, Y., Hu, H., Li, W., Li, S., Yang, J., Zhao, Q., & Zhang, C. (2023, June). Flow to control: Offline reinforcement learning with lossless primitive discovery. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 37, No. 9, pp. 10843-10851).

Questions:

1. What if the model learns inaccurately in a complex environment?

The inaccuracy of value functions and dynamics models will potentially lead the stitching process to a sub-optimal state, thus resulting in a failure during task execution. In our work, we adopt model ensemble to reduce the inaccuracy of both value functions and dynamics models.

2. Can you use the normalized score for the experimental results?

We use the absolute scores (the number of tasks accomplished) previously since we study the problem of task-level skill stitching. Normalized scores are just $\frac{\text{absolute scores}}{\text{the number of tasks in total}} \times 100\\%$. We have modified all the experimental results into normalized scores, as shown in the tables with green color in the revised paper.