Zero-Shot Offline Imitation Learning via Optimal Transport

Question 1

The paper is motivated by the fact that methods using goal-conditioned policies for zero-shot IL fail (Proposition 1 is used to support this argument). However, ZILOT plans using all the goal states in the demonstration. What if the goal-conditioned policy is given the expert demonstation as context? This way the goal-condition policy might discard the bad states to achieve a goal (shown in Proposition 1). Moreover, this would lead to learning a Markovian policy.

Thank you for raising this very interesting point. We agree that this represents a promising algorithmic idea. However, to the best of our knowledge, an RL algorithm that is conditioned on an entire trajectory would not be trivial. The machinery developed for goal-conditioned RL does not easily generalize to this setting (e.g., how would we sample “goal trajectories” for relabeling?). We believe that a naive implementation would then only be capable of imitating trajectories that appear in the offline dataset. Synthesising a (Markovian) reward function for multiple goals that cannot be “exploited” by repeatedly traversing through one of the goals also does not seem trivial. Our method represents a solution that sidesteps these issues by leveraging zero-order optimization, but we share the opinion that further research in zero-shot, model-free imitation to RL is also an interesting direction.

Question 2

The goodness of the proposed method depends on the planning horizon, and the paper discusses that ZILOT can be myopic too without a long planning horizon. This limits the applicability of the method in many real world tasks when planning for a long horizon. Is there a way to mitigate this by estimating the visitations using some form of bootstrapping?

We agree that bootstrapping represents a very interesting avenue of further research: in particular, estimating successor measures and future visitations would alleviate the finite horizon issue. We have decided to leave it for future research, as the current framework does not require an explicit policy representation, which would be necessary to estimate off-policy visitations, to the best of our knowledge. Moreover, we have observed empirically that myopic behavior is acceptable in practice, even for finite horizons.

Question 3

Since the method uses optimal transport and cites OTR [2] in the Related work, I feel it should be added as another baseline that uses the expert states and the sub-optimal dataset to recover a reward function and train a policy over it.

While OTR uses similar ideas to ours and was an inspiration for our work, it is not zero-shot and would require full training for each task we test as you have pointed out. For this reason we decided to not compare against this method. Finally, OTR is not straightforwardly applicable since it assumes access to the full expert states (as opposed to goals in our case), so a goal-conditioned value function would need to be trained additionally.

Question 4

Since sub-optimal data is used to learn the value-functions, can optimal V, W be recovered? If not, it is not highlighted how will this gap affect the final performance?

Since we are using an offline off-policy method to learn the goal-conditioned value function $V$, assuming convergence of the algorithm and a sufficient amount of data and coverage, the optimal $V^*$ (and thus also $W$) should be recovered. Of course, if the trained $V$ is far from optimal, ZILOT might either not take a short available path to a goal, or try to reach a goal in a way that is not possible (eg. through walls in maze).

Question 5

The abstract talks about zero-shot IL with existing practical methods viewing demonstrations as a sequence of goals. I feel this is not true as FB-IL [1] does zero-shot IL from a single demonstration. Moreover, I feel FB-IL should be a baseline. Although, ZILOT deals with demonstrations with a subset of states, I feel using Eq 8 in [1] should recover the policy with partial demonstrations.

Thank you for bringing this important line of research to our attention. We believe that this comparison highlights the novelty and the properties of ZILOT. Please refer to the general response for a detailed discussion.

We hope that we were able to clarify our contributions and address the reviewer’s comments. We remain open for further discussion.

We thank the reviewer for the detailed feedback, and for raising several interesting points of discussion. We have significantly updated our work based on the suggestions, and we believe that we can fully address all comments as follows.

Weakness 1

Many design choices of the proposed algorithm are not clear. Why does ZILOT use a non-Markovian policy as I believe the expert policies are Markovian?

We thank the reviewer for raising this important point, which we are happy to clarify. In standard RL, actions are selected to optimize an objective (i.e., the Q-function), which is conditioned on a specific policy. In this case, a Markovian policy can imitate a Markovian policy. ZILOT does not explicitly parameterize the policy, but computes actions in an online fashion, as is standard for zero-order optimization with learned world models [1]. Thus, the objective for action selection cannot be conditioned on a policy. Intuitively, when optimizing the action at timestep $t$, ZILOT needs to keep track of how the previous states align with the expert demonstration, in order to move towards states that are yet unvisited. In other words, this difference stems from the fact that ZILOT searches over actions, rather than over policies. Fortunately, this framework comes with an added benefit: ZILOT is fully able to imitate more general, non-Markovian policies.

[1] Hafner et. al. Learning Latent Dynamics for Planning from Pixels. ICML ‘19

Weakness 2

How is the function that maps states to goals defined / learned?

As established practice in the goal-conditioned RL literature [1], we consider the function $\phi$ as part of the problem definition. In our environments it captures the part of the state we want to control, e.g. the cube position in Fetch, the ball position in Pointmaze, and the x-axis and orientation of the Cheetah. The exact function is detailed in table 7 in the appendix.

[1] Andrychowicz et. al. Hindsight Experience Replay. NeurIPS ‘17

Weakness 3

It is not clear why the task success rate or the average returns are not used for comparing methods.

Our expert trajectories are not retrieved from optimizing a reward signal, but specified by humans: as rewards are undefined, we cannot evaluate returns. We can instead evaluate success rates, and this is what we indirectly report as GoalFraction. Compared to raw success rates (1 if all goals are achieved, else 0), we find GoalFraction (a scalar from 0 to 1, proportional to the number of goals in the expert demonstrations that are achieved in order) to be more fine-grained and informative.

Weakness 4

The limitations of the methods are not discussed. The only limitation presented is in Sec 7 that describes the dependence on a learned dynamics model.

Thank you for pointing this out, we have expanded limitations to a dedicated paragraph in Section 7.

Thank you for taking the time to review our work and your insightful comments. We have improved our paper based on your concerns, as addressed in the following.

Weakness

However, it would have been nice to see comparisons to other approaches to zero-shot imitation (e.g. along the line of FB representations)

Thank you for pointing out this important line of work. We analyse this family of methods in detail, compare it with ZILOT, and provide empirical evidence in Appendix B. Please refer to the general response for further discussion.

Question 1

How "in-distribution" are the inference tasks compared to the trajectories used to train TD-MPC? How does the quality of learned policy change we try to imitate more "OOD" expert behaviors?

We have added a paragraph in appendix C.8 investigating this question. We confirm that expert demonstrations do not need to appear in the offline dataset used to train TD-MPC2. However, it is important that individual states in the demonstrations are largely within the support of the dataset, even if they appear in separate trajectories. This can be interpreted as a form of off-policy learning, stitching. In other words, ZILOT can imitate “OOD” sequences of states, as long as each state is in-distribution.

We hope that we were able to clarify our contributions and address the reviewer’s comments. We remain open for further discussion.

We thank the reviewer for their feedback and comments. We are happy to address them, one by one, in this response.

Weakness 1

Lack of explanation and analysis of the myopic behavior in goal-conditioned imitation learning and why ZILOT bypass those limitations.

The fundamental issue with hierarchical approaches is that they “break down” the imitation learning problem into a sequence of local goal-conditioned problems. Each of the problems is then solved without considering the global imitation problem. As a consequence, these approaches can display greedy or myopic behavior, by aggressively trying to achieve the first subgoal, while disregarding the rest. ZILOT does not perform this decomposition in local problems, but tries to directly match the entire expert distribution (up to approximations due to finite-horizon planning). In other words, the policy extracted by ZILOT considers not only the next subgoal, but also the following ones. The suboptimality of hierarchical approaches is formally proved in Proposition 1, and can be observed directly in Figure 4. In Fetch-Slide, the naive baseline starts by only focusing on the first subgoal, and has to later correct the trajectory to accommodate for the following subgoals. By considering all subgoals, ZILOT produces a smoother trajectory that follows the expert more closely. We hope this discussion clarifies this important difference; we are happy to apply any further suggestion to improve our paper’s clarity.

Weakness 2

Baselines are relatively simple. How does it compare with some diffusion-based methods, e.g. Diffusion Policy?

To the best of our knowledge, diffusion-based methods for behavioral cloning need access to the expert’s actions, while our setting only provides sequences of expert states (or goals). The offline dataset $\mathcal{D}_\beta$ contains actions, but it does not demonstrate expert behavior: cloning this behavior would not recover an optimal goal-reaching policy. On the other hand, our baselines can handle suboptimal data, as they also build upon TD-MPC2. Since TD-MPC2 seems to be currently one of the best off-policy algorithms for online and offline RL [1], we are confident in the performance of our baselines. Nonetheless, in the context of this rebuttal, we evaluated an additional baseline (FB-IL): please see the common response for further details.

[1] Hansen et. al. D-MPC2: Scalable, Robust World Models for Continuous Control. ICLR ‘24

Question 1

Instead of conditioning on an intermediate goal, can we condition on a sequence of goals to mitigate myopic behavior?

Thank you for raising this interesting point. Given our analysis in Section 3, conditioning a policy on multiple future goals represents a promising direction for solving the issues faced by myopic policies. However, to the best of our knowledge there is no principled framework that trains such a policy from suboptimal data. This problem seems to be rather hard; even synthesising a reward function for multiple goals that cannot be exploited by repeatedly traversing through one of the goals does not seem to be trivial. Our method represents a first step into this direction.

Question 2

Instead of optimal transport to minimize the distance between distributions, how does it compare to diffusion-based methods?

While both diffusion models and optimal transport are concerned with distribution matching, their use cases differ: Optimal Transport, in particular Sinkorn’s algorithm that we use, is concerned with measuring the distance between two arbitrary distributions (over a certain space). Conversely, diffusion models are a solution to model one (usually very complex) distribution by minimizing a distribution matching objective. To the best of our knowledge, there is no established way of using a diffusion model to determine a distance between two distributions.

Thank you for your response! I appreciate the additional baseline (FB-IL) and most of my concerns are addressed. I will raise my score to 8.

Thank you for taking the time to review our work and your valuable feedback. We have improved our paper based on your comments, which we address in the following response.

Weakness

The proposed method requires both the learned world model and expert demonstrations. In practice, the training world model requires a lot of resources and either access to a simulator or a large collection of datasets.

it is still far away from “fair” zero-shot

We would like to clarify our setting and the requirements of the algorithm. As the reviewer points out, our method requires access to an offline dataset of trajectories (of arbitrary quality). However, ZILOT does not require a trained transition model as an input: through the experiments, the transition model is simply trained offline on the aforementioned dataset. Thus, at training time, ZILOT only requires access to an offline dataset. At test-time, ZILOT uses a single expert demonstration, only for specifying the task to be performed (i.e., it is not used for training). To the best of our knowledge, zero-shot imitation would not be possible without either of these two components. This notion of zero-shot imitation is consistent with the literature ([1]).

[1] Pathak et. al. Zero-Shot Visual Imitation. ICLR ‘18

Question 1

What is the additional computational cost introduced by solving OT problems compared to running just the planners?

For the fetch environments, the MPC+Cls planner runs at around 30Hz compared to ZILOT which runs at around 3Hz. We should point out that there exists a range of runtime improvements on the basic Sinkhorn Algorithm that we did not test for this work. We included these details in Appendix C.3.

Question 2a

How does the method scale when there are more than 1 demonstration?

Thank you for suggesting this interesting setting. As ZILOT deals with occupancies, it can directly generalize to imitation of multiple demonstrations, by including all given trajectories in the empirical approximation of $\rho^E$ and truncating each trajectory individually for the finite-horizon problem that is optimized with MPC. Notably, the same approach can also be applied to the occupancy estimate of the agent policy $\rho^\pi$ which allows ZILOT to build upon probabilistic world models, eg. diffusion-based world-models [1].

[1] Alonso et. al. Diffusion for World Modeling: Visual Details Matter in Atari. NeurIPS ’24

Question 2b

Does the method have a problem with multimodal behavior in the bigger data collection?

ZILOT relies on offline learning for both its dynamics model and value function. For the former, we adopt the standard RSSM ([1]) in TD-MPC2, while the latter is learned through an offline reinforcement learning algorithm. Both algorithms are fully capable of handling multimodal behavior. In fact, the offline dataset $\mathcal{D}_\beta$ is highly multi-modal; nevertheless we do not observe significant issues during training.

[1] Shaj et. al. Hidden Parameter Recurrent State Space Models for changing Dynamics Scenarios. ICLR ‘22

Question 3 & 4

Can the authors apply the method to cross-domain / cross-embodiment settings? What if the expert policy is different from the resulting policy in TD-MPC? I.E can we observe a distributional shift if our expert policy comes from an environment with slightly different dynamics, for example, cross-embodiment (x-magical https://github.com/kevinzakka/x-magical)?

As long as the goal-spaces coincide, the expert can have a different embodiment as ZILOT imitates experts from rough (potentially incomplete) and partial (goal-space only) trajectories and does not need expert actions. Further, we would also like to point out that TD-MPC2, the method used in ZILOT to train its dynamics model and value functions, shows great performance in cross-domain/cross-embodiment settings ([1]). We are happy to further discuss this points in case it remains unclear.

[1] Hansen et. al. D-MPC2: Scalable, Robust World Models for Continuous Control. ICLR ‘24

(continued)