Zero-Shot Offline Imitation Learning via Optimal Transport

## Abstract

Abstract Zero-shot imitation learning algorithms hold the promise of reproducing unseen behavior from as little as a single demonstration at test time.
Existing practical approaches view the expert demonstration as a sequence of goals, enabling imitation with a high-level goal selector, and a low-level goal-conditioned policy.
However, this framework can suffer from myopic behavior: the agent’s immediate actions towards achieving individual goals may undermine long-term objectives.
We introduce a novel method that mitigates this issue by directly optimizing the occupancy matching objective that is intrinsic to imitation learning.
We propose to lift a goal-conditioned value function to a distance between occupancies, which are in turn approximated via a learned world model.
The resulting method can learn from offline, suboptimal data, and is capable of non-myopic, zero-shot imitation, as we demonstrate in complex, continuous benchmarks.
Code for this project is provided at https://github.com/martius-lab/zilot .

## 1 Introduction

The emergence of zero/few-shot capabilities in language modeling (Brown et al., ; Wei et al., ; Kojima et al., ) has renewed interest in generalist agents across all fields in machine learning.
Typically, such agents are pretrained with minimal human supervision.
At inference, they are capable of generalization across diverse tasks, without further training, i.e. zero-shot.
Such capabilities have also been a long-standing goal in learning-based control (Duan et al., ).
Promising results have been achieved by leveraging the scaling and generalization properties of supervised learning (Jang et al., ; Reed et al., ; O’Neill et al., ; Ghosh et al., ; Kim et al., ), which however rely on large amounts of expert data, usually involving costly human participation, e.g. teleoperation.
A potential solution to this issue can be found in reinforcement learning approaches, which enable learning from suboptimal data sources (Sutton & Barto, ).
Existing methods within this framework ease the burden of learning general policies by limiting the task class to additive rewards (Laskin et al., ; Sancaktar et al., ; Frans et al., ) or single goals (Bagatella & Martius, ).

This work lifts the restriction of previous approaches, and proposes a method that can reproduce rich behaviors from offline, suboptimal data sources.
In particular, we allow arbitrary tasks to be specified through a single demonstration at inference time, conforming to a zero-shot Imitation Learning (IL) framework.
From a practical standpoint, this demonstration may be partial (i.e., lack action labels) and rough (e.g., only contain a small set of abstract key states to be reached).
For example, when tasking a robot arm with moving an object along a path, it is sufficient to provide the object’s position for a few “checkpoints” without specifying the exact pose that the arm has when each checkpoint is reached.

In principle, a specified goal sequence can be decomposed into multiple single-goal tasks that can be accomplished by goal-conditioned policies, as proposed by recent zero-shot IL approaches (Pathak et al., ; Hao et al., ).
However, we show that this decomposition is prone to myopic behavior.
Continuing the robotic manipulation example from above, let us consider a task described by two sequential goals, each specifying a certain position that the object should reach.
In this case an optimal goal-conditioned policy would attempt to reach the first goal as fast as possible, and possibly throw the object towards it.
The agent would then relinquish control of the object, leaving it in a suboptimal—or even unrecoverable—state.
In this case, the agent would be unable to move the object towards the second goal.
This myopic behavior is a fundamental issue arising from goal abstraction, as we formally argue in Section 3, and results in catastrophic failures in hard-to-control environments, as we demonstrate empirically in Section 5.

In this work we instead provide an holistic solution to zero-shot offline imitation learning by adopting an occupancy matching formulation.
We name our method ZILOT (Zero-shot Offline Imitation Learning from Optimal Transport).
We utilize Optimal Transport (OT) to lift the state-goal distance inherent to GC-RL to a distance between the expert’s and the policy’s occupancies, where the latter is approximated by querying a learned world model.
Furthermore, we operationalize this distance as an objective in a standard fixed horizon MPC setting.
Minimizing this distance leads to non-myopic behavior in zero-shot imitation.
We verify our claims empirically by comparing our planner to previous zero-shot IL approaches across multiple robotic simulation environments, down-stream tasks, and offline datasets.

## 2 Preliminaries

### 2.1 Imitation learning

We model an environment as a controllable Markov Chain(^1^11or reward-free Markov Decision Process.) ${\mathcal{M}}=({\mathcal{S}},{\mathcal{A}},P,\mu_{0})$, where ${\mathcal{S}}$ and ${\mathcal{A}}$ are state and action spaces, $P:{\mathcal{S}}\times{\mathcal{A}}\to\Omega({\mathcal{S}})$(^2^22where $\Omega({\mathcal{S}})$ denotes the set of distributions over ${\mathcal{S}}$.) is the transition function and $\mu_{0}\in\Omega({\mathcal{S}})$ is the initial state distribution.
In order to allow for partial demonstrations, we additionally define a goal space ${\mathcal{G}}$ and a surjective function $\phi:{\mathcal{S}}\to{\mathcal{G}}$ which maps each state to its abstract representation.
To define “goal achievement”, we assume the existence of a goal metric $h$ on ${\mathcal{G}}$ that does not need to be known.
We then regard state $s\in{\mathcal{S}}$ as having achieved goal $g\in{\mathcal{G}}$ if we have $h(\phi(s),g)<\epsilon$ for some fixed $\epsilon>0$.
For each policy $\pi:{\mathcal{S}}\to\Omega({\mathcal{A}})$, we can measure the (undiscounted) $N$-step state and goal occupancies respectively as

$$ $\displaystyle\varrho_{N}^{\pi}(s)=\frac{1}{N+1}\sum_{t=0}^{N}\Pr[s=s_{t}]\quad \text{and}\quad\rho_{N}^{\pi}(g)=\frac{1}{N+1}\sum_{t=0}^{N}\Pr[g=\phi(s_{t})],$ (1) $$

where $s_{0}\sim\mu_{0}$, $s_{t+1}\sim P(s_{t},a_{t})$ and $a_{t}\sim\pi(s_{t})$.
These quantities are particularly important in the context of imitation learning.
We refer the reader to Liu et al. () for a full overview over IL settings, and limit this discussion to offline IL.
Specifically, we assume access to two datasets:
${\mathcal{D}}_{\beta}=(s_{0}^{i},a_{0}^{i},s_{1}^{i},a_{1}^{i},\dots)_{1}^{|{
\mathcal{D}}_{\beta}|}$ consisting of full state-action trajectories from $\mathcal{M}$ and ${\mathcal{D}}_{E}=(g_{0}^{i},g_{1}^{i},\dots)_{1}^{|{\mathcal{D}}_{E}|}$ containing demonstrations of an expert in the form of goal sequences, not necessarily abiding to the dynamic of $\mathcal{M}$.
Note that both datasets do not have reward labels.
The goal is to train a policy $\pi$ that imitates the expert, which is commonly formulated as matching goal occupancies

$$ $\displaystyle\rho_{N}^{\pi}\overset{D}{=}\rho_{N}^{\pi_{E}}.$ (2) $$

The only additional constraint imposed by zero-shot offline IL is that ${\mathcal{D}}_{E}$ consists of just one goal-sequence $(g_{0},\dots,g_{M})=g_{0:M}$, and is only available at inference time.

### 2.2 Optimal Transport

In the field of machine learning, it is often of interest to match distributions, i.e. find some probability measure $\mu$ that resembles some other probability measure $\nu$.

Figure: Figure 2: An example of Optimal Transport between the discrete approximation $\hat{\mu},\hat{\nu}$ of two Gaussians $\mu,\nu$. The cost matrix $C$ consists of the point-wise costs where the cost here is the Euclidian distance. A coupling matrix $T\in\mathcal{U}(\hat{\mu},\hat{\nu})$ (middle) is visualized through lines representing the matching (right).
Refer to caption: x2.png

In recent years there has been an increased interest in Optimal Transportation (OT) (Amos et al., ; Haldar et al., ; Bunne et al., ; Pooladian et al., ).
As illustrated in figure 2, OT does not only compare probability measures in a point-wise fashion, like $f$-Divergences such as the Kullbach-Leibler Divergence ($D_{\text{KL}}$), but also incorporates the geometry of the underlying space.
This also makes OT robust to empirical approximation (sampling) of probability measures (Peyré & Cuturi (), p.129).

Formally, OT describes the coupling $\gamma\in{\mathcal{P}}({\mathcal{X}}\times{\mathcal{Y}})$ of two measures $\mu\in{\mathcal{P}}({\mathcal{X}}),\nu\in{\mathcal{P}}({\mathcal{Y}})$ with minimal transportation cost w.r.t. some cost function $c:{\mathcal{X}}\times{\mathcal{Y}}\rightarrow{\mathbb{R}}$.
The primal Kantorovich form is given as the optimization problem

$$ $\displaystyle\text{OT}_{c}(\mu,\nu)=\inf_{\gamma\in{\mathcal{U}}(\mu,\nu)}\int _{{\mathcal{X}}\times{\mathcal{Y}}}c(x_{1},x_{2})d\gamma(x_{1},x_{2})$ (3) $$

where the optimization is over all joint distributions of $\mu$ and $\nu$ denoted as $\gamma\in{\mathcal{U}}(\mu,\nu)$ (couplings).
If ${\mathcal{X}}={\mathcal{Y}}$ and $({\mathcal{X}},c)$ is a metric space
then for $p\in{\mathbb{N}}$, $W_{p}^{p}=\text{OT}_{c^{p}}$ is called the Wasserstein-$p$ distance which was shown to be a metric on the subset of measures on ${\mathcal{X}}$ with finite $p$-th moments (Clement & Desch, ).

Given samples $x_{1},\dots,x_{n}\sim\mu$ and $y_{1},\dots,y_{m}\sim\nu$ the discrete OT problem between the discrete probability measures $\hat{\mu}=\sum_{i=1}^{n}{a}_{i}\delta_{x_{i}}$ and $\hat{\nu}=\sum_{j=1}^{m}{b}_{j}\delta_{y_{j}}$
can be written as a discrete version of equation 3, namely

$$ $\displaystyle\text{OT}_{c}(\hat{\mu},\hat{\nu})=\min_{{\bm{T}}\in\mathcal{U}({ \bm{a}},{\bm{b}})}\sum_{i=1}^{n}\sum_{j=1}^{m}c(x_{i},y_{j}){T}_{ij}=\min_{{ \bm{T}}\in\mathcal{U}({\bm{a}},{\bm{b}})}\langle{\bm{C}},{\bm{T}}\rangle$ (4) $$

with the cost matrix ${C}_{ij}=c(x_{i},y_{j})$.
The marginal constraints can now be written as
$\mathcal{U}({\bm{a}},{\bm{b}})=\{{\bm{T}}\in{\mathbb{R}}^{n\times m}:{\bm{T}}
\cdot\bm{1}_{m}={\bm{b}}\text{ and }{\bm{T}}^{\top}\cdot\bm{1}_{n}={\bm{a}}\}$.
This optimization problem can be solved via Linear Programming.
Furthermore, Cuturi () shows that the entropically regularized version, commonly given as $\text{OT}_{c,\eta}(\hat{\mu},\hat{\nu})=\min_{{\bm{T}}\in\mathcal{U}({\bm{a}},
{\bm{b}})}\langle{\bm{C}},{\bm{T}}\rangle-\eta D_{\text{KL}}({\bm{T}},{\bm{a}}
{\bm{b}}^{\top})$, can be efficiently solved in its dual form using Sinkhorn’s algorithm (Sinkhorn & Knopp, ).

### 2.3 Goal-conditioned reinforcement learning

As techniques from the literature will be recurring in this work, we provide a short introduction to fundamental ideas in GC-RL.
We can introduce this framework by enriching the controllable Markov Chain ${\mathcal{M}}$. We condition it on a goal $g\in{\mathcal{G}}$ and cast it as an (undiscounted) Markov Decision Process ${\mathcal{M}}_{g}=({\mathcal{S}}\cup\{\perp\},{\mathcal{A}},P_{g},\mu_{0},R_{g
},T_{\max})$. Compared to the reward-free setting above, the dynamics now include a sink-state $\perp$ upon goal-reaching and a reward of $-1$ until this happens:

$$ $\displaystyle P_{g}(s,a)=\begin{cases}P(s,a)&\text{if }h(\phi(s),g)\geq \epsilon\\ \delta_{\perp}&\text{otherwise}\end{cases},\ R_{g}(s,a)=\begin{cases}-1&\text{ if }h(\phi(s),g)\geq\epsilon\\ 0&\text{otherwise}\end{cases}$ (5) $$

where $\delta_{x}$ stands for the probability distribution assigning all probability mass to $x$.

We can now define the goal-conditioned value function as

$$ $\displaystyle V^{\pi}(s_{0},g)=\mathop{\mathbb{E}}_{\mu_{0},P_{g},\pi}\left[ \sum_{t=0}^{T_{\max}}R_{g}(s_{t},a_{t})\right]\text{ where }s_{0}\sim\mu_{0},s _{t+1}\sim P_{g}(s_{t},a_{t}),a_{t}\sim\pi(s_{t},g).$ (6) $$

The optimal goal-conditioned policy is then $\pi^{\star}=\operatorname*{arg\,max}_{\pi}\mathbb{E}_{g\sim\mu_{{\mathcal{G}}}
,s\sim\mu_{0}}V^{\pi}(s_{0};g)$ for some goal distribution $\mu_{{\mathcal{G}}}\in\Omega({\mathcal{G}})$.
Intuitively, the value function $V^{\pi}(s,g)$ corresponds to the negative number of expected steps that $\pi$ needs to move from state $s$ to goal $g$.
Thus the distance $d=-V^{\star}$ corresponds to the expected first hit time.
If no goal abstraction is present, i.e. $\phi=\text{id}_{\mathcal{S}}$, then $({\mathcal{S}},d)$ is a quasimetric space (Wang et al., ), i.e. $d$ is non-negative and satisfies the triangle inequality.
Note, though, that $d$ does not need be be symmetric.

## 3 Goal abstraction and myopic planning

The distribution matching objective at the core of IL problems is in general hard to optimize.
For this reason, most practical methods for zero-shot IL leverage a hierarchical decomposition into a sequence of GC-RL problems (Pathak et al., ; Hao et al., ).
We will first describe this approach, and then show how it potentially introduces myopic behavior and suboptimality.

In the pretraining phase, Pathak et al. () propose to train a goal-conditioned policy $\pi_{g}:{\mathcal{S}}\times{\mathcal{G}}\to{\mathcal{A}}$ on reaching single goals and a goal-recognizer $C:{\mathcal{S}}\times{\mathcal{G}}\to\{0,1\}$ that detects whether a given state achieves the given goal.
Given an expert demonstration $g_{1:M}$ and an initial state $s_{0}$, imitating the expert can then be sequentially decomposed into $M$ goal-reaching problems, and solved with a hierarchical agent consisting of two policies.
On the lower level, $\pi_{g}$ chooses actions to reach the current goal; on the higher level, $C$ decides whether the current goal is achieved and $\pi_{g}$ should target the next goal in the sequence.

We define the pre-image $\phi^{-1}(g)=\{s\in{\mathcal{S}}:\phi(s)=g\}$ as the set of all states that map to a goal, and formalize the suboptimality of the above method under goal abstraction as follows.

Figure: Figure 3: Controllable Markov Chain with $\phi:(x,y)\mapsto x$.

###### Proposition 1 .

Let us define the optimal classifier $C(s,g)=\mathbf{1}_{h(\phi(s),g)<\epsilon}$.
Given a set of visited states $\mathcal{P}\subseteq\mathcal{S}$, the current state $s\in\mathcal{P}$, and a goal sequence $g_{1:M}\in\mathcal{G}^{M}$, let the optimal hierarchical policy be $\pi_{h}^{\star}(s)=\pi^{\star}(s,g_{i+1})$, where $i$ is the smallest integer such that there exist a state $s_{p}\in\mathcal{P}$ with $h(\phi(s_{p}),g_{i})<\epsilon$, and $i=0$ otherwise.
There exists a controllable Markov Chain $\mathcal{M}$ and a realizable sequence of goals $g_{0:M}$ such that, under a suitable goal abstraction $\phi(\cdot)$, $\pi^{\star}_{h}$ will not reach all goals in the sequence, i.e. $\rho^{\pi^{\star}_{h}}_{N}(g_{i})=0$ for some $i\in[0,\dots,M]$ and all $N\in\mathbb{N}$.

###### Proof.

Consider the Markov Chain ${\mathcal{M}}$ depicted in figure 3 with goal abstraction $\phi:(x,y)\mapsto x$ and $p>0$.
Now, consider the goal sequence $(g_{0},g_{1},g_{2})=(0,1,2)$, which can only be achieved, by a policy taking action $a_{1}$ in the initial state $s_{0}=(0,0)$.
Consider $\pi^{\star}_{h}$ in $s_{0}$, with $\mathcal{P}=\{s_{0}\}$.
The smallest integer $i$ such that $h(\phi(s_{0}),g_{i})<\epsilon$ is $i=0$, therefore $\pi_{h}^{\star}(s_{0})=\pi^{\star}(s_{0},g_{1})$.
We can then compare the state-action values $Q$ in $s_{0}$:

$$ $\displaystyle Q^{\pi^{\star}(\cdot,g_{1})}(s_{0},a_{1},g_{1})=\sum_{t=0}^{T_{ \max}}-p^{t}=-1\cdot\frac{1-p^{(T_{\max}+1)}}{1-p}\ <\ -1=Q^{\pi^{\star}(\cdot ,g_{1})}(s_{0},a_{0},g_{1}).$ (7) $$

This implies that $\pi_{h}^{\star}(s_{0})=\pi^{\star}(s_{0},1)=a_{0}$.
The next state visited by $\pi_{h}^{\star}$ will always be $(1,0)$, from which $(2,1)$ is not reachable, and $g_{2}$ is not achievable. We thus have $\rho^{\pi^{\star}_{h}}_{N}(g_{2})=0$ for all $N\in\mathbb{N}$.
∎

We remark that this issue arises in the presence of goal abstraction which plays a vital role in the partial demonstration setting we consider.
Without goal abstraction, i.e., if each goal is fully specified, there is no leeway in how to achieve it for the policy (assuming $\epsilon\to 0$ as well).
Nevertheless, goal abstraction is ubiquitous in practice (Schaul et al., ) and necessary to enable learning in complex environments (Andrychowicz et al., ).

###### Proposition 1 .

###### Proof.

## 4 Optimal Transport for Zero-Shot IL

Armed with recent tools in value estimation, model-based RL and trajectory optimization, we propose a method for zero-shot offline imitation learning that directly optimizes the occupancy matching objective, introducing only minimal approximations.
As a result, the degree of myopia is greatly reduced, as we show empirically in section 5.

In particular, we propose to solve the occupancy matching problem in equation 2 by minimizing the Wasserstein-$1$ metric $W_{1}$ with respect to goal metric $h$ on the goal space ${\mathcal{G}}$, i.e.

$$ $\displaystyle W_{1}(\rho_{N}^{\pi},\rho_{N}^{E})=\text{OT}_{h}(\rho_{N}^{\pi}, \rho_{N}^{E}).$ (8) $$

This objective involves two inaccessible quantities: goal occupancies $\rho_{N}^{\pi},\rho_{N}^{E}$, as well as the goal metric $h$.
Our key contribution lies in how these quantities can be practically estimated, enabling optimization of the objective with scalable deep RL techniques.

##### Occupancy Estimation

Since the expert’s and the policy’s occupancy are both inaccessible, we opt for discrete, sample-based approximations.
In the case of the expert occupancy $\rho_{N}^{E}$, the single trajectory provided at inference $(g_{0},\dots,g_{M})$ represents a valid sample from it, and we use it directly.
For an arbitrary agent policy $\pi$, we use a discrete approximation after training a dynamics model $\hat{P}\approx P$ on ${\mathcal{D}}_{\beta}$, which can be done offline through standard supervised learning.
We can then approximate $\rho_{N}^{\pi}$ by jointly rolling out the learned dynamics model and the policy $\pi$.
We thus get the discrete approximations

$$ $\displaystyle\rho_{N}^{E}\approx\hat{\rho}_{M}^{E}=\frac{1}{M+1}\sum_{j=0}^{M} \delta_{g_{j}}\ \text{ and }\ \rho_{N}^{\pi}\approx\hat{\rho}_{N}^{\pi}=\frac{ 1}{N+1}\sum_{t=0}^{N}\delta_{\phi({s_{t}})}$ (9) $$

where for the latter we sample $s_{0}\sim\mu_{0},s_{t+1}\sim\hat{P}(s_{t},a_{t}),a_{t}\sim\pi(s_{t})$.
Similarly, we can also obtain an estimate for the state occupancy of $\pi$ as $\varrho_{N}^{\pi}\approx\hat{\varrho}_{N}^{\pi}=\frac{1}{N+1}\sum_{t=0}^{N}
\delta_{s_{t}}$.

##### Metric Approximation

As $h$ may be unavailable or hard to specify in practical settings, we propose to train a goal-conditioned value function $V^{\star}$ from the offline data $\mathcal{D}_{\beta}$ and use the distance $d(s,g)=-V^{\star}(s,g)$ (i.e. the learned first hit time) as a proxy.
For a given state-goal pair $(s,g)$, this corresponds to the approximation $d(s,g)\approx h(\phi(s),g)$.
It is easy to show that a minimizer of $h(\phi(\cdot),g)$ also minimizes $d(\cdot,g)$.
Using $d$ also has the benefit of incorporating the dynamics of the MDP into the cost of the OT problem.
The use of this distance has seen some use as the cost function in Wasserstein metrics between state occupancies in the past (Durugkar et al., ).
As we show in section 5.3, $d$ is able to capture potential asymmetries in the MDP, while remaining informative of $h$.
We note that, while $h:{\mathcal{G}}\times{\mathcal{G}}\to\mathbb{R}$ is a distance in goal-space, $d:{\mathcal{S}}\times{\mathcal{G}}\to\mathbb{R}$ is a distance between states and goals. Nonetheless, $d$ remains applicable as the policy’s occupancy can also be estimated in state spaces as $\hat{\varrho}_{N}^{\pi}$.
Given the above considerations, we can rewrite our objective as the discrete optimal transport problem

$$ $\displaystyle\pi^{\star}=\operatorname*{arg\,min}_{\pi}\text{OT}_{d}(\hat{ \varrho}_{N}^{\pi},\hat{\rho}_{M}^{E}).$ (10) $$

##### Optimization

Having addressed density and metric approximations, we now focus on optimizing the objective in equation 10.
Fortunately, as a discrete OT problem, the objective can be evaluated efficiently using Sinkhorn’s algorithm when introducing entropic regularization with a factor $\eta$ (Cuturi, ; Peyré & Cuturi, ).
A non-Markovian, deterministic policy optimizing the objective at state $s_{k}\in\mathcal{S}$ can be written as

$$ $\displaystyle\pi(s_{0:k},g_{0:m})\approx\operatorname*{arg\,min}_{a_{k}}\min_{ a_{k+1:N-1}}\text{OT}_{d,\eta}\left(\frac{1}{N+1}\sum_{i=0}^{N}\delta_{s_{i}}, \frac{1}{M+1}\sum_{j=0}^{M}\delta_{g_{j}}\right)$ (11) $$

where $s_{0:k}$ are the states visited so far and $s_{k+1:N}$ are rolled out using the learned dynamics model $\hat{P}$ and actions $a_{k:N-1}$.
Note that while $s_{0:k}$ are part of the objective, they are constant and are not actively optimized.

Intuitively, this optimization problem corresponds to finding the first action from a sequence ($a_{k:N-1}$) that minimizes the OT costs between the empirical expert goal occupancy, and the induced empirical policy state occupancy.
This type of optimization problem fits naturally into the framework of planning with zero-order optimizers and learned world models (Chua et al., ; Ha & Schmidhuber, ); while these algorithms are traditionally used for additive costs, the flexibility of zero-order optimizers (Rubinstein & Kroese, ; Williams et al., ; Pinneri et al., ) allows a straightforward application to our problem. The objective in equation 11 can thus be directly optimized with CEM variants (Pinneri et al., ) or MPPI (Williams et al., ), in a model predictive control (MPC) fashion.

Like for other MPC approaches, we are forced to plan for a finite horizon $H$, which might be smaller than $N$, because of imperfections in the learned dynamics model or computational constraints.
This is referred to as receding horizon control (Datko, ).
When the policy rollouts used for computing $\hat{\varrho}^{\pi}_{N}$ are truncated, it is also necessary to truncate the goal sequence to exclude any goals that cannot be reached within $H$ steps.
To this end, we train an extra value function $W$ that estimates the number of steps required to go from one goal to the next by regressing onto $V$, i.e. by minimizing $\mathbb{E}_{s,s^{\prime}\sim{\mathcal{D}}_{\beta}}[(W(\phi(s);\phi(s^{\prime})
)-V(s;\phi(s^{\prime})))^{2}]$.
For $i\in[0,\dots,M]$, we can then estimate the time when $g_{i}$ should be reached as

$$ $\displaystyle t_{i}\approx-V(s_{0};g_{0})-\sum_{j=1}^{i}W(g_{j-1};g_{j}).$ (12) $$

We then simply truncate the online problem to only consider goals relevant to $s_{1},\dots,s_{k+H}$, i.e. $g_{0},\dots,g_{K}$ where $K=\min\{j:t_{j}\geq k+H\}$.
We note that this approximation of the infinite horizon objective can potentially result in myopic behavior if $K<M$; nonetheless, optimal behavior is recovered as the effective planning horizon increases.

Algorithm 1 shows how the practical OT objective is computed.

Figure: Algorithm 1 OT cost computation for ZILOT

##### Implementation

The method presented relies solely on three learned components: a dynamics model $\hat{P}$, and the state-goal and goal-goal GC value functions $V$ and $W$.
All of them can be learned offline from the dataset ${\mathcal{D}}_{\beta}$.
In practice, we found that several existing deep reinforcement learning frameworks can be easily adapted to learn these functions.
We adopt TD-MPC2 (Hansen et al., ), a state of the art model-based algorithm that has shown promising results in single- and multitask online and offline RL.
We note that planning takes place in the latent space constructed by TD-MPC2’s encoders.
We adapt the method to allow estimation of goal-conditioned value functions, as described in appendix B.
We follow prior work (Andrychowicz et al., ; Bagatella & Martius, ; Tian et al., ) and sample goals from the future part of trajectories in ${\mathcal{D}}_{\beta}$ in order to synthesize rewards without supervision.
We note that this goal-sampling method also does not require any knowledge of $h$.

## 5 Experiments

This section constitutes an extensive empirical evaluation of ZILOT for zero-shot IL.
We first describe our experimental settings in terms of environment, baselines and metrics, and then present qualitative and quantitative result, as well as an ablation study.
We consider a selection of 30 tasks defined over 5 environments, as summarized below and described in detail in appendix A.

fetch (Plappert et al., ) is a manipulation suite in which a robot arm either pushes (Push), or lifts (Pick&Place) a cube towards a goal.
We adopt these two environments directly.
To illustrate the failure cases of myopic planning, we also evaluate a variation of Push (i.e. Slide), in which the table size exceeds the arm’s range, the table’s friction is reduced, and the arm is constrained to be always touching the table.
As a result, the agent cannot fully constrain the cube, e.g. by picking it up, or pressing on it, and the environment strongly punishes careless manipulation.
In all three environments, tasks consist of moving the cube along trajectories shaped like the letters “S”, “L”, and “U”.

halfcheetah (Wawrzyński, ) is a classic Mujoco environment where the agent controls a cat-like agent in a 2D horizontal plane.
As this environment is not goal-conditioned by default, we choose the x-coordinate and the orientation of the cheetah as a meaninful goal-abstraction.
This allows the definition of tasks involving standing up and hopping on front or back legs, as well as doing flips.

pointmaze (Fu et al., ) involves maneuvering a pointmass through a maze via force control. Downstream tasks consist of following a series of waypoints through the maze.

##### Planners

The most natural comparison is the framework proposed by Pathak et al. (), which addresses imitation through a hierarchical decomposition, as discussed in section 3.
Both hierarchical components are learned within TD-MPC2: the low-level goal-conditioned policy is by default part of TD-MPC2, while the goal-classifier (Cls) can be obtained by thresholding the learned value function $V$.
We privilege this baseline (Policy+Cls) by selecting the threshold minimizing $W_{\min}$ per environment among the values $[1,2,\dots,5]$.
Moreover, we also compare to a version of this baseline replacing the low-level policy with zero-order optimization of the goal-conditioned value function (MPC+Cls), thus ablating any benefits resulting from model-based components.
We remark that all MPC methods use the same zero-order optimizer iCEM (Pinneri et al., ).

##### Metrics

We report two metrics for evaluating planner performance.
The first one is the minimal encountered (empirical) Wasserstein-$1$ Distance under the goal metric $h$ of the agent’s trajectory and the given goal sequence.
Formally, given trajectory $(s_{0},\dots,s_{N})$ and the goal sequence $(g_{0},\dots,g_{M})$ we define

$$ $\displaystyle W_{\min}(s_{0:N},g_{1:M}):=\min_{k\in\{0,\dots,N\}}W_{1}\left( \frac{1}{k+1}\sum_{i=0}^{k}\delta_{\phi(s_{i})},\frac{1}{M+1}\sum_{j=0}^{M} \delta_{g_{j}}\right).$ (13) $$

This metric takes the minimum over the trajectory length as it is in general hard to estimate the exact number of steps needed to imitate a goal sequence.
We introduce a secondary metric “GoalFraction” since $W_{\text{min}}$ does not evaluate the order in which goals are reached.
It represents the fraction of goals that are achieved in the order they were given.
Formally, this corresponds to the length of the longest subsequence of achieved goals that matches the desired order.

### 5.1 Can ZILOT effectively imitate unseen trajectories?

We first set out to qualitatively evaluate whether the method is capable of imitation in complex environments, despite practical approximations. Figure 4 illustrates how Pi+Cls, MPC+Cls, and ZILOT imitate an expert sliding a cube across the big table of the fetch_slide_large_2D environment.
Both myopic baselines struggle to regain control over the cube after moving it towards the second goal, leading to straight trajectories that leave the manipulation range.
In contrast, ZILOT plans beyond the second goal.
As displayed in the middle part of figure 4, the coupling of the OT problem approximately pairs up each state in the planned trajectory with the appropriate goal.
This leads to closer imitation of the expert, as shown in the renders.

Figure: Figure 4: Example tasks in fetch_slide_large_2D. The left three columns show five trajectories across five seeds of both myopic methods we evaluate (Pi+Cls, MPC+Cls) and ZILOT (ours). The trajectories are drawn in the $x$-$y$-plane of the goal space and just show the movement of the cube. ZILOT’s behavior imitates the given goal trajectories more closely. On the right, we visualize the OT objective at around three quarters of the episode time. It includes both the past and planned future states, as well as their coupling to the goals. Note that planning occurs in the latent state of TD-MPC2, and separately trained decoders are used for this visualization.
Refer to caption: x3.png

### 5.2 How does ZILOT perform compared to prior methods?

We provide a quantitative evaluation of ZILOT with respect to myopic methods in table 1. For more details we refer the reader to appendix A.
As ZILOT directly optimizes a distribution matching objective, it generally reproduces expert trajectories more closely, achieving a lower Wasserstein distance to its distribution.
This is especially evident in environments that are very punishing to myopic planning, such as the Fetch Slide environment shown in figure 4.
In most environments, our method also out-performs the baselines in terms of the fraction of goals reached.
In less punishing environments, ZILOT may sacrifice precision in achieving the next goal exactly for an overall closer match of the expert trajectory.
This is most clearly visible in the pointmaze environment.
We note that the performance of the two baselines is comparable to each other’s, suggesting that the performance gap to ZILOT stems from the change in objective, rather than implementation or model-based components.

Table: Table 1: Performance of Pi+Cls, MPC+Cls and ZILOT (ours) in all environments and tasks. Each metric is the mean over 20 trials, we then report the mean and standard deviation of those metrics across 5 seeds. We perform a Welch $t$-test with $p=0.05$ do distinguish the best values and mark them bold. Values are rounded to 3 and 2 digits respectively.

### 5.3 What matters for ZILOT?

To validate some of our design choices we finally evaluate the following versions of our method.

- •
OT+unbalanced, our method with unbalanced OT (Liero et al., ; Séjourné et al., ), which turns the hard marginal constraint $\mathcal{U}$ (see section 2.2) into a soft constraint. We use this method to address the fact that a rough expert trajectory may not necessarily yield a feasible expert occupancy approximation.
- •
OT+Cls, a version of our method which includes the goal-classifier (Cls), with the same hyperparameter search performed for the baselines. This method discards all past states and goals that are recognized as reached, and does not consider them when computing and matching occupancies.
- •
OT+$h$, our method with the goal metric $h$ on ${\mathcal{G}}$ as the cost function in the OT problem, replacing $d$.

Our results are summarized in figure 5.
First, we see that using unbalanced OT does not yield significant improvements.
Second, using a goal-classifier can have a bad impact on matching performance.
We suspect this is the case because keeping track of the history of states gives a better, more informative, estimate of which part of the expert occupancy has already been fulfilled.
Finally, we observe that the goal metric $h$ may not be preferable to $d$, even if it is available.
We mainly attribute this to the fact that, in the considered environments, any action directly changes the state occupancy, but the same cannot be said for the goal occupancy.
Since $h$ only allows for the comparison of goal occupancies, the optimization landscape can be very flat in situations where most actions do not change the future state trajectory under goal abstraction, such as the start of fetch tasks as visible in its achieved trajectories in the figures in appendix C.
Furthermore, while $h$ is locally accurate, it ignores the global geometry of MDPs, as shown by its poor performance in strongly asymmetric environments (i.e., halfcheetah).

Figure: Figure 5: Ablation of design choices in ZILOT, including coupling constraints (OT+unbalanced), partial trajectory matching (OT+Cls), and the approximation of $h$ by $d$ (OT+$h$). For more detailed results, please refer to table 2.
Refer to caption: x4.png

## 6 Related Work

##### Zero-shot IL

When a substantial amount of compute is allowed at inference time, several methods have been proposed to leverage pretrained models to infer actions, and retrieve an imitator policy via behavior cloning (Pan et al., ; Zhang et al., ; Torabi et al., ). As already discussed in section 3, most (truly) zero-shot methods cast the problem of imitating an expert demonstration as following the sequence of its observations (Pathak et al., ; Hao et al., ).
Expert demonstrations are then imitated by going from one goal to the next using a goal-conditioned policy.
This approach can also be extended to a setting where only the final goal is considered absolutely necessary and intermediate ones can be filtered out if they are suboptimal w.r.t. reaching the final goal (Hao et al., ).
In contrast, our work proposes a holistic approach to imitation, which considers all goals within the planning horizon.

##### Zero-Shot RL

Vast amounts of effort have been dedicated to learning generalist agents without supervision, both on the theoretical (Touati & Ollivier, ; Touati et al., ) and practical side (Laskin et al., ; Mendonca et al., ).
Among others, (Sancaktar et al., ; P. et al., ; Bagatella & Martius, ) learn a dynamics model through curious exploration and show how it can be leveraged to optimize additive objectives.
More recently, Frans et al. () use Functional Reward Encodings to encode arbitrary additive reward functions in a latent that is used to condition a policy.
While these approaches are effective in a standard RL setting, they are not suitable to solve instances of global RL problems (Santi et al., ) (i.e., distribution matching).

##### Imitation Learning

A range of recent work has been focused on training agents that imitate experts from their trajectories by matching state, state-action, or state-next-state occupancies depending on what is available.
These methods either directly optimize various distribution matching objectives (Liu et al., ; Ma et al., ) or recover a reward using Generative Adversarial Networks (GAN) (Ho & Ermon, ; Li et al., ) or in one instance OT (Luo et al., ).
Another line of work has shown impressive real-world results by matching the action distributions (Shafiullah et al., ; Florence et al., ; Chi et al., ) directly.
All these approaches do not operate in a zero-shot fashion, or need ad-hoc data collection.

##### OT in RL

Various previous work has used Optimal Transport in RL as a reward signal.
One application is online fine-tuning where a policy’s rollouts are rewarded in proportion to how closely they match expert trajectories or the rollouts of experts (Dadashi et al., ; Haldar et al., ).
Luo et al. () instead use a similar trajectory matching strategy to recover reward labels for unlabelled mixed-quality offline datasets.
Most of the works mentioned above do not have any special metric or cost-function they use for their OT problems.
The most common choices are Cosine Similarities and Euclidean distances for their general applicability.

## 7 Conclusion

In this work, we point out a failure-mode of current zero-shot IL methods that cast imitating an expert demonstration as following a sequence of goals with myopic GC-RL policies.
We address this issue by framing the problem as occupancy matching.
By introducing discretizations and minimal approximations, we derive an Optimal Transportation problem that can be directly optimized at inference time using a learned dynamics model, goal-conditioned value functions, and zero-order optimizer. Our experimental results across various environments and tasks show that our approach outperforms state-of-the-art zero-shot IL methods, particularly in scenarios where non-myopic planning is crucial.
We additionally validate our design choices through a series of ablations.

The main practical limitation of our implementation is the reliance on a learned dynamics model, which may be inaccurate over long horizons.
This forces the optimization of a fixed-horizon objective, which reintroduces a slight degree of myopia, as the agent may fail to consider goals beyond the planning horizon.
However, we found the degree of myopia to be acceptable in our experimental settings, and expect our framework to become more and more applicable as the accuracy of learned world models improves.

#### Acknowledgments

We thank Anselm Paulus, Mikel Zhobro, and Núria Armengol Urpí for their help throughout the project.
Marco Bagatella and Jonas Frey are supported by the Max Planck ETH Center for Learning Systems.
Georg Martius is a member of the Machine Learning Cluster of Excellence, EXC number 2064/1 – Project number 390727645.
We acknowledge the support from the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039B).

## Appendix A Additional Results

### A.1 Main Result Details

In table 2 we provide detailed results for all ablations.
We also provide a summarized version of the results in figure 6.

Figure: Figure 6: Summarized performance of all discussed Planners. See table 1 and table 2 for detailed results.
Refer to caption: x5.png

Table: Table 2: Performance of our method and its ablations in all environments and tasks. Each metric is the mean over 20 trials, we then report the mean and standard deviation of those metrics across 5 seeds. We perform a Welch $t$-test with $p=0.05$ do distinguish the best values and mark them bold. Values are rounded to 3 and 2 digits respectively.

### A.2 Finite Horizon Ablations

As discussed in section 4, we are forced to optimize the objective over a finite horizon $H$ due to the imperfections in the learned dynamics model and computational constraints.
The hyperparameter $H$ should thus be as large as possible, as long as the model remains accurate.
We visualize this trade-off in figure 7 for environment fetch_slide_large_2D.
It is clearly visible that if the horizon is smaller than 16, the value we chose for our experiments, then performance rapidly deteriorates towards the one of the myopic planners.
However, when increasing the horizon beyond 16, performance does not improve, suggesting that the model is not accurate enough to plan beyond this horizon.

Figure: Figure 7: Mean performance across five seeds in fetch_slide_large_2D for different planning horizons.
Refer to caption: x6.png

### A.3 Single Goal Performance

When the expert trajectory consists of only a single goal, myopic planning is of course sufficient to imitate the expert.
To verify this we evaluate the performance of all planners in the standard single goal task of the environments.
Figure 8 shows the success rate of all planners in this task verifying that non-myopic planning neither hinders nor helps in this case.

Figure: Figure 8: Single Goal Success Rate in the standard single goal tasks of the environments. We report the mean performance across 20 trials and standard deviation across 5 seeds.
Refer to caption: x7.png

## Appendix B Implementation Details

### B.1 ZILOT

The proposed method is motivated and explained in section 4.
We now present additional details.

##### Sinkhorn

First, we rescale the matrix ${\bm{C}}$ by $T_{\max}$ and clamp it to the range $[0,1]$ before running Sinkhorns algorithm.
The precise operation performed is

$$ $\displaystyle{\bm{C}}\leftarrow\min\left(1,\max(0,{\bm{C}}/T_{\max})\right).$ (14) $$

This is done so that the same entropy regularization $\epsilon$ can be used across all environments, and to ensure there are no outliers that hinder the convergence of the Sinkhorn algorithm.
For the algorithm itself, we use a custom implementation for batched OT computation, heavily inspired by Flamary et al. () and Cuturi et al. ().
We run our Sinkhorn algorithm for $r=500$ iterations with a regularization factor of $\epsilon=0.02$.

##### Truncation

When the agent gets close to the end of the expert trajectory, then we might have that $t_{K}<k+H$, i.e. the horizon is larger than needed.
We thus truncate the planning horizon to the estimated remaining number of steps (and at least $1$), i.e. we set

$$ $\displaystyle H_{\text{actual}}\leftarrow\max\left(1,\min(t_{K}-k,H)\right).$ (15) $$

##### Unbalanced OT

As mentioned in the main text in section 5.3, we can use unbalanced OT (Liero et al., ; Séjourné et al., ) to address that fact that the uniform marginal for the goal occupancy approximation may not be feasible.
Unbalanced OT replaces this hard constraint of ${\bm{T}}\top\cdot\bm{1}_{N}=\bm{1}_{M}$ into the term $\xi_{b}\text{KL}({\bm{T}}\top\cdot\bm{1}_{N},\bm{1}_{M})$ in the objective function.
For our experiments we have chosen $\xi_{b}=1$.

### B.2 TD-MPC2 Modifications

As TD-MPC2 (Hansen et al., ) is already a multi-task algorithm that is conditioned on a learned task embedding $t$ from a task id $i$, we only have to switch out this conditioning to a goal latent $z_{g}$ to arrive at a goal-conditioned algorithm as detailed in table 3.
We remove the conditioning on the encoders and the dynamics model $f$ completely as the goal conditioning of GC-RL only changes the reward but not the underlying Markov Decision Process ${\mathcal{M}}$ (assuming truncation after goal reaching, see section 2.3).
For training we adopt all TD-MPC2 hyperparameters directly (see table 8).
As mentioned in the main text, we also train a small MLP to predict $W$ that regresses on $V$.

**Table 3: Our modifications to TD-MPC2 to making it goal- instead of task-conditioned.**
|  | TD-MPC2 (Hansen et al., ) | “GC”-TD-MPC2 (our changes) |
| --- | --- | --- |
| Task/Goal Embedding | $t=E(i)$ | $z_{g}=h_{g}(g)$ |
| Encoder | $z=h(s,t)$ | $z=h(s)$ |
| Dynamics | $z^{\prime}=f(z,a,t)$ | $z^{\prime}=f(z,a)$ |
| Reward Prediction | $r=R(z,a,t)$ | $r=R(z,a,z_{g})$ |
| $Q$-function | $q=Q(z,a,t)$ | $q=Q(z,a,z_{g})$ |
| Policy | $a\sim\pi(z,t)$ | $a\sim\pi(z,z_{g})$ |

We have found the computation of pair-wise distances $d$ to be the major computational bottleneck in our method, as TD-MPC2 computes them as $d=-V^{\pi}(s,g)=-Q(z,\pi(z,z_{g}),z_{g})$ where $z=h(s),z_{g}=h_{g}(g)$.
To speed-up computation, we train a separate network that estimates the value function directly.
It employs a two-stream architecture (Schaul et al., ; Eysenbach et al., ) of the form $V^{\pi}(z,z_{g})=\phi(z)^{\top}\psi(z_{g})$ where $\phi$ and $\psi$ are small MLPs for fast inference of pair-wise distances.
Using this network architecture for $V$, ZILOT runs at 0.5 to 3Hz on an Nvidia GTX 2080ti GPU, depending on the size of $H$ and the size of the OT problem.
Several further steps could be taken to speed-up the sinkhorn algorithm itself, including $\eta$-schedules and/or Anderson acceleration (Cuturi et al., ) as well as warm-starting it with potentials, e.g. from previous (optimizer) steps or with a trained network (Amos et al., ).

### B.3 Goal Sampling

Figure: Table 4: Goal Sampling

As mentioned in the main text, we follow prior work (Andrychowicz et al., ; Bagatella & Martius, ; Tian et al., ) and sample goals from the future part of trajectories in ${\mathcal{D}}_{\beta}$ in order to synthesize rewards without supervision.
The exact procedure is as follows:

- •
With probability $p_{\text{future}}$ we sample a goal from the future part of the trajectory with time offset $t_{\Delta}\sim\text{Geom}(1-\gamma)$.
- •
With probability $p_{\text{next}}$ we sample the next goal in the trajectory.
- •
With probability $p_{\text{rand}}$ we sample a random goal from the dataset.

See table 4 for the hyperparameters used.

### B.4 Training

We train our version of TD-MPC2 offline with the datasets detailed in table 5 for 600k steps.
Training took about 8 to 9 hours on a single Nvidia A100 GPU.
Note that as TD-MPC2 samples batches of 3 transitions per element, we effectively sample $3\cdot 256=768$ transitions per batch.
The resulting models are then used for all planners and experiments.

Table: Table 5: Environment description. We detail the datasets used for training.

### B.5 Environments

We detail environment details in table 6.
Note that while we consider an undiscounted setting, we specify $\gamma$ for the goal sampling procedure above.

Table: Table 6: Environment details. We detail the goal abstraction $\phi$, metric $h$, threshold $\epsilon$, horizon $H$, maximum episode length $T_{\max}$, and discount factor $\gamma$ used for each environment.

The environments fetch_push and fetch_pick_and_place and pointmaze_medium are used as is.
As halfcheetah is not goal-conditioned by default, we define our own goal range to be $(x,\theta_{y})\in[-5,5]\times[-4\pi,4\pi]$(^3^33Note that the halfcheetah environment does not reduce $\theta$ with any kind of modular operation, i.e. states with $\theta=0$ and $\theta=2\pi$ are distinct.).
fetch_slide_large_2D is a variation of the fetch_slide environment where the table size exceeds the arm’s range and the arm is restricted to two-dimensional movement touching the table.

### B.6 Tasks

The tasks for the fetch and pointmaze environments are specified in the environments normal goal-space. Their shapes can be seen in the figures in appendix C.
To make the tasks for halfcheetah more clear, we visualize some executions of our method in the figures 9, 10, 11, 12, 13, and 14.

Figure: Figure 9: Example trajectory of ZILOT (ours) in halfcheetah-backflip-running.
Refer to caption: extracted/5916754/graphics/halfcheetah_vis/backflip-running/2/vis.png

Figure: Figure 10: Example trajectory of ZILOT (ours) in halfcheetah-backflip.
Refer to caption: extracted/5916754/graphics/halfcheetah_vis/backflip/0/vis.png

Figure: Figure 11: Example trajectory of ZILOT (ours) in halfcheetah-frontflip-running.
Refer to caption: extracted/5916754/graphics/halfcheetah_vis/frontflip-running/0/vis.png

Figure: Figure 12: Example trajectory of ZILOT (ours) in halfcheetah-frontflip.
Refer to caption: extracted/5916754/graphics/halfcheetah_vis/frontflip/0/vis.png

Figure: Figure 13: Example trajectory of ZILOT (ours) in halfcheetah-hop-backward.
Refer to caption: extracted/5916754/graphics/halfcheetah_vis/hop-backward/0/vis.png

Figure: Figure 14: Example trajectory of ZILOT (ours) in halfcheetah-hop-forward.
Refer to caption: extracted/5916754/graphics/halfcheetah_vis/hop-forward/3/vis.png

### B.7 Hyperparameters

**Table 7: Hyperparameters used for iCEM (Pinneri et al., ). We use the implementation from Pineda et al. ().**
| Name | Value |
| --- | --- |
| num_iterations | 4 |
| population_size | 512 |
| elite_ratio | 0.01 |
| population_decay_factor | 1.0 |
| colored_noise_exponent | 2.0 |
| keep_elite_frac | 1.0 |
| alpha | 0.1 |

**Table 8: TD-MPC2 Hyperparameters. We have adopted these unchanged from Hansen et al. ()**
| Name | Value |
| --- | --- |
| lr | 3e-4 |
| batch_size | 256 |
| n_steps (“horizon”) | 3 |
| rho | 0.5 |
| grad_clip_norm | 20 |
| enc_lr_scale | 0.3 |
| value_coef | 0.1 |
| reward_coef | 0.1 |
| consistency_coef | 20 |
| tau | 0.01 |
| log_std_min | -10 |
| log_std_max | 2 |
| entropy_coef | 1e-4 |

## Appendix C Additional Qualitative Results

In the following, we present all goal-space trajectories across all planners, tasks, and seeds presented in this work.
Note that since the tasks of the fetch environments display some natural symmetries, we decided to split evaluations between all four symmetrical versions of them.
Further, we quickly want to stress that these trajectories are shown in goal-space.
This means that if the cube in fetch is not touched, as is the case in some cases for ZILOT+$h$, then the trajectory essentially becomes a single dot at the starting position.
Also note that Pi+Cls is completely deterministic, which is why its visualization appears to have less trajectories.

Figure: (a) U-dense
Refer to caption: x8.png

Figure: (a) L-sparse
Refer to caption: x16.png

Figure: (a) S-dense
Refer to caption: x24.png

Figure: (a) U-dense
Refer to caption: x32.png

Figure: (a) L-sparse
Refer to caption: x40.png

Figure: (a) S-dense
Refer to caption: x48.png

Figure: (a) U-dense
Refer to caption: x56.png

Figure: (a) L-sparse

Figure: (a) S-dense
Refer to caption: x72.png

Figure: (a) circle-sparse
Refer to caption: x80.png

Figure: (a) run-forward
Refer to caption: x84.png

## Appendix D Goal Classifier Hyperparameter Search

As mentioned in the main text, we perform an extensive hyperparameter search for the threshold value of the goal classifier (Cls) for the myopic methods Pi+Cls and MPC+Cls as well as for the ablation of our method ZILOT+Cls.
In figures 27 and 26 we show the performance of the three respective planners in all five environments and denote the threshold values that yield the best performance per environment.
Interestingly, in some of the fetch environments not all tasks attain maximum performance with the same threshold value showing that this hyperparameter is rather hard to tune.

Figure: (a) fetch_pick_and_place
Refer to caption: x92.png

Figure: (a) Pi+Cls for fetch_pick_and_place
Refer to caption: x97.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. What is the additional computational cost introduced by solving OT problems compared to running just the planners?
2. How does the method scale when there are more than 1 demonstration? Does the method have a problem with multimodal behavior in the bigger data collection?
3. Can the authors apply the method to cross-domain / cross-embodiment settings?
4. 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)?
5. Can the current method be altered to remove the model-based approach and replace it with the offline dataset of video collection? I.e why can't we learn value functions from just video data without rewards and actions and then use it in OT?
soundness: 3
strengths: The method shows example of non-myopic behaviour compared to the baselines. The proposed approach does not depened on availability of actions in the expert trajectory. The proposed approach works as zero-shot "agent correction" policy. Authors evaluated their method on variety of different environments.
summary: The authors introduce a new method for zero-shot imitation learning based on optimal transport.  
The approach involves combining a modified goal-conditioned TD-MPC2 algorithm with optimal transport. The authors begin by training goal-conditioned value functions V and W, as well as a dynamics model obtained from TD-MPC2. Then they run the OT process using the Sinkhorn algorithm, given V, W, P, and the expert trajectory, to compute the cost matrix and transport matrix.
weaknesses: 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.

Though, the authors positioned their paper as "zero-shot IL" it is still far away from "fair" zero-shot since it requires some conditions for the method to work:
1. Access to the offline dataset 
2. Trained Transition model
3. Expert trajectory.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: - Instead of conditioning on an intermediate goal, can we condition on a sequence of goals to mitigate myopic behavior?
- Instead of optimal transport to minimize the distance between distributions, how does it compare to diffusion-based methods?
soundness: 3
strengths: - Identifies an existing problem of goal-conditioned policy learning and proposes a solution based on optimal transport.
- Demonstrate success of ZILOT on a range of locomotion and manipulation tasks and show better alignment with demonstrations.
summary: This paper identifies an existing problem with goal-conditioned imitation learning: an agent that optimized to achieve individual goals can undermine long-term objectives. Instead of learning a goal-conditioned policy, ZILOT learns a dynamics model with a goal-conditioned value function. It uses optimal transport to minimize the divergence between the rollout policy and expert demonstration.

Experiments show that ZILOT imitates demonstration more closely compared with MPC and goal-conditioned policy.
weaknesses: - Lack of explanation and analysis of the myopic behavior in goal-conditioned imitation learning and why ZILOT bypass those limitations.
- Baselines are relatively simple. How does it compare with some diffusion-based methods, e.g. Diffusion Policy?

---

confidence: 2
contribution: 4
initial_rating: 6
presentation: 4
questions: 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?
soundness: 3
strengths: The paper studies a very important problem: the ability to quickly learn new skills and policies from few expert demonstrations is a much-desired property of policies. The method is conceptually simple and well-motivated theoretically.

I found the paper to be easy to read, and the background for all relevant concepts was motivated and well-introduced (within Sec 2-4)

Figure 4 clearly demonstrates the main improvement over the sub-goal approach, where the MPC + CIs overshoot the puck when going to the first subgoal.

The experiments are relatively thorough (both in the main paper, and the Appendix), indicating improvement over other model-based zero-shot imitation methods.
summary: This paper proposes an approach for zero-shot imitation -- use a single expert trajectory to extract a policy to mimic this task. The general approach is to solve an optimal transport problem (equipped by a goal-reaching cost function) with a model that can be used to estimate the policy's stationary distribution. Across a number of tasks (fetch, halfcheetah, pointmaze), the proposed approach can learn useful skills.
weaknesses: I had no major objections to the paper; 

However, it would have been nice to see comparisons to other approaches to zero-shot imitation (e.g. along the line of FB representations. Another axis that could improve the paper would be to see a wider range of downstream tasks (requiring the synthesis of more diverse motions),  perhaps those from (Frans et al)

Minor nit: There is a lot of content in the Appendix, but it is not clearly linked within the main paper. I would encourage the authors to link and discuss this content within the main paper, so that it is not missed by a reader.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: 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.
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? 
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. 
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? 
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.

### References
[1] Pirotta et al., Fast Imitation via Behavior Foundation Models. ICLR'24\
[2] Luo et al., Optimal transport for offline imitation learning, ICLR'23
soundness: 2
strengths: - The idea of recovering expert behaviors from partial trajectories with subgoals and without action labels is interesting and challenging. 
- The drawback of prior methods that used goal-conditioned policies for this task is discussed well to motivate the proposed method.
summary: This paper presents a method for zero-shot imitation learning with a single demonstration (without action labels or partial trajectories consisting of a sequence of goals). A method called ZILOT that learns to plan and act according to the goals is proposed. Firstly, ZILOT learns a dynamics model (similar to TD-MPC2) using a dataset of transitions which can be sub-optimal. For planning using the demonstration, a non-Markovian method is employed to match the occupancy of the partial demonstration and the state-visitation distribution of the policy. The discrepancy between occupancies is computed using Optimal Transport which requires value functions: V to get reachability of a goal state and W to get the steps taken between goals. The non-Markovian planner uses this discrepancy to select the best action. Experiments conducted on multiple benchmarks show that ZILOT is better at W1 distance between expert and policy occupancies and better at generating trajectories that follow the order of goals in the demonstration.
weaknesses: - 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? How is the function $\phi$ that maps states to goals defined / learned?
- The experiments are not convincing as the baselines like FB-IL [1] and OTR [2] are missing. It is not clear why the task success rate or the average returns are not used for comparing methods. 
- 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.

Meta-review: This paper proposes ZILOT (Zero-shot Offline Imitation Learning via Optimal Transport), a novel method for zero-shot imitation learning that aims to address the myopic behavior issue present in existing approaches that decompose expert demonstrations into sequences of individual goals and target each goal independently. The key innovation is directly optimizing an occupancy matching objective using optimal transport between the expert's and policy's state distributions, rather than treating goals independently. This allows ZILOT to consider multiple future goals during planning, leading to more sophisticated non-myopic behavior. The empirical results across manipulation and locomotion tasks show that ZILOT outperforms baseline methods, particularly in scenarios requiring longer-term planning. The reviewers appreciated the novelty and motivation of the approach as well as the theoretical underpinnings of the approach.
The main weaknesses are: (a) ZILOT requires solving computationally expensive optimal transport problems during planning, leading to significantly slower inference compared to baselines, (b) the finite planning horizon means some degree of myopic behavior remains, though less severe than baselines, (c) missing comparisons to relevant baselines like FB-IL and OTR, though these were partially addressed in the rebuttal, (d) questions about whether the added complexity of non-Markovian policies is truly necessary when Markovian policies may be sufficient. Overall, the paper had mixed set of reviews and during the internal discussion phase, reviewer Lfwv reiterated some of the above limitations which the authors are encouraged to address in a subsequent version.

Decision: Reject