Definition of zero-shot

We would like to clarify our definition of “zero-shot”. We refer to a method as “zero-shot” if it retrieves an optimal policy for unseen objectives provided at test-time, with modest compute overhead. “Zero-shot” methods may be allowed a compute-heavy pre-training phase, which should not be informed of the downstream task. This definition aligns with [1], which is thus a method for zero-shot IL. Similarly, [2] and [3] perform zero-shot IL and RL, respectively. IQLearn, TAILO, PW-DICE, SMODICE, AILOT, and the majority of offline IL methods, would not qualify as zero-shot IL methods, as they need the objective (in the case of IL, the demonstration) to be available during pre-training.

Existing definitions do not specify what constitutes a “modest compute overhead”. In [1], “zero-shot” methods require <5 minutes to imitate a single trajectory, while non-zero-shot methods require >3 hours (fig. 2 in [1]). ZILOT takes <4 minutes to imitate a trajectory in our tasks, which is comparable to more involved FB-IL variants, and orders of magnitude faster than non-zero-shot methods, e.g., BC. We would thus argue that ZILOT performs zero-shot IL.

This definition of “zero-shot” depends on when demonstrations are provided, not on the information they contain. They may contain actions (as for FB-IL-BC), or may not. We can qualify this distinction: we may define a method to be “learning from observations” if it imitates demonstrations that do not contain actions. ZILOT and [2] would then belong to this category. For both, an action-labeled offline dataset (not demonstrations, but arbitrary trajectories) is provided during pre-training to convey action semantics.

We thank the reviewer for encouraging this discussion, which we hope clarifies our definitions. We will update our submission accordingly.

Concerns - Finite sample approximation of OT

We acknowledge that our occupancy matching objective is sample-based and thus subject to approximation errors. As mentioned, for both OTR and our work, this approach performs well empirically. We hypothesize that this might be due to the adoption of entropy regularization [4], or to a low dimensionality (as defined in [5]) of occupancy under smooth dynamics in discrete time.

Suggested Improvements

World Model Training

We will expand Appendix C.2 with training objectives for our practical choice of world model (TD-MPC2), and refer to it prominently in the main text.

Goal Abstraction

We introduce the general concept of a goal-abstraction in Sec. 2.1. and specify the exact functions used for each environment we evaluate in Table 5 in the Appendix (e.g., for Fetch we use the cube pose). As is standard in GC-RL the goal-abstraction is known [6]. ZILOT’s planning objective (Eq. 17) matches the sampled state occupancy to the expert goal occupancy, and thus uses a standard GCVF $V(s, g)$.

Eq. 18 is used for estimating the time the expert required for the demonstration; as the demonstration only contains goals, it requires a separate value function estimating distances between pairs of goals ($W(g, g’)$). It is only used as a heuristic for selecting a part of the expert trajectory to match at each step, which is necessary for finite-horizon optimization.

Extra Evaluation Environments

Please refer to our response to Reviewer LRC1, where we motivate our choice of environments and provide evaluations of the Walker environment.

Other Weaknesses - Running Time

We would like to clarify that the reported planning frequencies were recorded on dated hardware (GTX 2080 TI). On modern hardware (RTX 4090) the planning frequencies are at least 2-4Hz depending on the problem size. We will update these numbers in the manuscript. Further, we expect that the inference time could easily be doubled with a more efficient implementation of the Sinkhorn Algorithm using JIT-compilation [7] or writing specialized kernels [8].

Question on Expert mismatch

As our expert demonstrations may be rough and partial, our method can imitate expert demonstrations sourced from a slightly different embodiment (e.g. the new evaluations on Walker reuse the expert demonstrations for Cheetah).

---

We hope this addresses your questions, and we are happy to elaborate on points that were inadequately covered due to character limitations.

References:

1. Pirotta et. al. Fast Imitation via Behavior Foundation Models. ICLR ‘24
2. Pathak et. al. Zero Shot Visual Imitation. ICLR ‘18
3. Touati et. al. Learning one Representation to optimize all Rewards. NeurIPS ‘21
4. Genevay et. al. Sample Complexity of Sinkhorn Convergences. AISTATS ‘19
5. Weed et al. Sharp Asymptotic and Finite-Sample Rates of Convergence of Empirical Measures in Wasserstein Distance.
6. Andrychowicz et. al. Hindsight Experience Replay. NeurIPS ‘17
7. Cuturi et. al. Optimal Transport Tools (OTT) arXiv:2201.12324
8. Feydy. et. al. Interpolating between Optimal Transport and MMD using Sinkhorn Divergences. AISTATS ‘19

Thank you for your assessment and valuable feedback. We first address your main concern, our empirical evaluation, and then your other questions.

Q5 on further experiments

Our experiments are chosen to be representative of the 3 most common types of MDPs found in robotics: manipulation (Fetch), navigation (Pointmaze), and locomotion (Cheetah). Instead of including similar environments, we focused on evaluating a diverse set of tasks in each environment (see Table 1), in order to cover the full complexity of possible behaviors. We have performed evaluations on additional Mujoco tasks, i.e. Walker with the same tasks as in Cheetah, and found results to be strongly consistent with Cheetah:

Our evaluation is also similar to prior work in zero-shot Imitation from offline data. [2] evaluate on different tasks in one manipulation environment. [1] uses four embodiments all of which are locomotion tasks. In comparison, our evaluation is richer.

Q1 model-free IL adaptation

To the best of our knowledge, zero-shot distribution matching under OT has not been explored with model-free methods. The model-based component allows us to accurately predict and optimize finite-horizon occupancies. A model-free variant optimizing a similar OT objective could use techniques from [1], which we found to be, however, less data efficient (Section 5.2). Thus, a model-free variant represents an interesting research direction.

Q2 on OT vs. other distances

The main advantages of OT over f-divergences (eg. KL) is that (1) it is more robust to empirical approximation (line 100 col 2), and (2) it can incorporate the underlying geometry of the space(s) which allows us to use an MDP-specific (learned) metric, the goal-conditioned (GC) value function in our case (line 210 col 2).

Q3 on finite horizon planning

We acknowledge that our method relies on finite horizon planning in practice; as we show empirically, this is sufficient to largely avoid myopic behavior. An extension would be possible leveraging occupancy estimates (e.g., as in [1]); however, these techniques did not perform well in our evaluations (Section 5.2).

Q4 on sampling errors

Estimation errors from sampling did not seem to impact performance in our experiments, especially as the support of the future state distribution is often modest. If the learned world model is non-deterministic, one may sample multiple trajectories from it to get a better estimate of the future state distribution.

Q6 on the comparison to IL and IRL

ansewredThe large majority of methods for imitation learning and inverse RL (e.g., GAIL, IQLearn, DICE, OTR) are not zero-shot, and require demonstrations to be provided in advance, at training time. After pre-training, ZILOT is instead capable of imitating unseen trajectories. To the best of our knowledge, only two model-free approaches can achieve the same. The first approach [2] is however fundamentally myopic (Section 3), and is considered as a baseline (Cls+Policy in table 1). The second approach [1] is model-free, zero-shot, and non-myopic; however, we find it to be less data efficient, and to underperform in our evaluation (Section 5.2). Thus, ZILOT distinguishes itself from existing model-free approaches because it is zero-shot, non-myopic, and data-efficient.

Q7 on scalability

ZILOT relies on a GC value function and a world model. Both of these are well-studied objects in the field, and we expect ZILOT to scale as these components improve.

---

We hope this fully addresses your questions and comments, and we are happy to expand on any answers that had to be short due to character limitations.

References:

1. Pirotta et. al. Fast Imitation via Behavior Foundation Models. ICLR ‘24
2. Pathak et. al. Zero Shot Visual Imitation. ICLR ‘18

We would like to thank the reviewer for their feedback and comments.

Thank you for pointing out that the acronym “GC-RL” was not introduced in the paper. We will introduce it properly given a chance to update the paper.

The subscripts in the definitions of $\mathcal{D}\_\beta$ and $\mathcal{D}\_E$ were meant to denote that $i$ runs from $1$ to $|\mathcal{D}\_\beta|$ and $|\mathcal{D}\_E|$ respectively. We will update the definitions to $\mathcal{D}\_\beta = (s\_0^i, a\_0^i, s\_1^i, a\_1^i, \dots)\_{i=1}^{|\mathcal{D}\_\beta|}$ and $\mathcal{D}\_E = (g\_0^i, g\_1^i, …)\_{i=1}^{|\mathcal{D}\_E|}$ respectively, to make this more clear.

We remain available for further discussion for the rest of the rebuttal period.