Zero-Shot Offline Imitation Learning via Optimal Transport
Authors: Thomas Rupf
Marco Bagatella
Nico Gürtler
Jonas Frey
Georg Martius
Venue: ICML 2025 poster (2025)
Keywords: Imitation Learning
Deep Reinforcement Learning
Optimal Transport
Primary area: reinforcement_learning
TL;DR: A non-myopic method for zero-shot imitation from arbitrary offline data.

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. The code is available at https://github.com/martius-lab/zilot.

Zero-Shot Offline Imitation Learning via Optimal Transport
Thomas Rupf! Marco Bagatella" 23
Nico Gürtler 12
Jonas Frey? 23 Georg Martius 12
Abstract
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 reinforce- ment learning approaches, which enable learning from suboptimal data sources (Sutton & Barto, 2018). Existing methods within this framework ease the burden of learning general policies by limiting the task class to additive rewards (Laskin et al., 2021: Sancaktar et al., 2022; Frans et al., 2024) single goals (Bagatella & Martius, 2023).
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 high-level goal selector, and a low-level goal-conditioned policy. However, this framework can suffer from myopic behavior: the agent's 's immediate actions towards achieving indi- vidual goals may undermine long-term objectives. We introduce a novel method that mitigates this issue by directly optimizing the occupancy match- ing objective that is intrinsic to imitation learning. We propose to lift conaiiondi value func- tion to distance between occupancies, which are in turn approximated via learned world model. The resulting method can learn from offline, sub- optimal data, and is capable of non- mon-mc, zero- shot imitation, as we demonstrate in complex, continuous benchmarks. The code is available at https:: //tp::///iie /ttp//gttttttt////////////ttehht zil poto..
This work lifts the restrictions of previous approaches, and proposes a method that can reproduce rich behaviors from offlin,, suboptimal data sources. We allow arbitrary tasks to be specified through a single demonstration devoid of actions at inference time, conforming to zero-shot Imitation Learning (IL) setting (Pathak et al., 2018; Pirotta et al., 2024). Furthermore, we consider a relaxation of this setting (Pathak et al., 2018), where the expert demonstration may be rough, consisting of an ordered sequence of states without precise time-step information, and partial, meaning each state contains only partial information about the full state. These two relaxations are desirable from a practical standpoint, as they allow user to avoid specifying infor- mation that is either inconsequential the task costly to attain (e.g. only through teleoperation). 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 arm pose.
1. Introduction
The emergence of zero/few-shot capabilities in language modeling (Brown et al., 2020; Wei et al., 2022; Kojima al., 2022) 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, with- out further training, i.c. zero-shot. Such capabilities have also been a long-standing goal in learning-base based control (Duan et al., 2017). Promising results have been achieved by leveraging the scaling and generalization properties of supervised learning (Jang et al., 2022; Reed al., 2022: O'Neill al., 2023; Ghosh et al., 2024; Kim et al., 2024),
In principle, a specified goal sequence can be decomposed into multiple single-goal tasks that can be accomplished by goal-conditioned policies, proposed by recent zero-shot IL approaches (Pathak et al., 2018; Hao et al., 2023). However, we show that this decomposition is prone to myopic behavior when expert demonstrations are partial. Continuing the robotic manipulation example from above, let us consider task described by two sequential goals, each specifying certain position that the object should reach. In this case optimal goal-conditioned policy would attempt to reach the first goal as fast as possible, possibly by throwing the object towards it. The agent would then relinquish control of the object, leaving it in subop- timal-or or even unrecoverable state. In this case, the agent would be unable to move the object towards the second goal. This myopic behavior is fundamental issue arising from goal abstraction,
'Universität Tübingen, Tübingen, Germany 'MP for Intelli- gent Systems, Tübingen, Germany 'ETH, Zürich, Switzerland. Correspondence to: Thomas Rupf <thrupf ethz.ch...
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
formally argue in Section

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Training
World Model
Inference
Zero-shot IL
P
oooce
Expert
Offline Dataset
Past
D
Future
TT
3
Rollout P
Goal
Minimize
E
Conditioned
V
OT-v,(# .
Myopic
Value Function
Non-Myopic (ours)
Figure 1. Overview of ZILOT. After learning world model P and condiconnt value function V from offline data (left), zero-order optimizer directly matches the occupancy of rollouts 6" from the learned world model the occupancy a single expert demonstration & (center). This is done by lifting the goal-conditioned value function distance between occupancies using Optimal Transport. The resulting policy displays non-myopic behavior (right).
3. and results in catastrophic failures in hard-to-control and environments, we demonstrate empirically in Section 5.
N Σ 1=0
In this work we instead provide a 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 dis- tance inherent GC-RL distance between the expert's and the policy's 's occupancies, where the latter is approxi- mated by querying learned world model. Furthermore, we operationalize this distance 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.
(T) N+1
Prig = ((x))}
(2)
where so Po. 8t+1 ~ P(st, and a, ~ (T()). These quantities are particularly important in the context of imita- tion learning. We refer the reader to (Liu et al., 2023) for full overview over IL settings, and limit this discussion to of- fline IL with rough and partial expert trajectories. It assumes access to datasets: D3 = (s a sf, ai ;=1 consisting of full state-action trajectories from M and
|aal
DE Si:ggi.... ) containing demonstrations of an expert in the form of goal sequences, not necessarily abid- ing to the dynamic of M. Note that both datasets do not have reward labels. The goal is to train policy r that imi- tates the expert, which is commonly formulated as matching goal occupancies
2. Preliminaries
D PNN
2.1. Imitation Learning
PN
We model an environment as a controllable Markov Chain = -(S, A. P, Po), where S and A are state and action spaces, P: x A (2S)) is the transition function and 140 E (2(S) the initial state distribution. In order to allow for partial demonstrations, we additionally define goal space G and a surjective function @ S + G which maps each state to its abstract representation. To define "goal achievement", we assume the existence of a goal metric on g that does not need to be known. We then regard state E as having achieved goal g E if we have For each policy
The setting consider in this work is zero-shot offline IL which imposes two additional constraints offline IL. First, Dr is only available at inference time, which means pre-training has to be task-agnostic. We further assume Dg consists of a single trajectory (gu..... = 90.m- Second, at inference, the agent should imitate nE. with a "modest compute-overheant (Pathak et al., 2018: Touati & Ollivier, 2021; Pirotta et al., 2024). In practice, imitation of unseen trajectories should be order of magnitudes cheaper than IL from scratch, and largely avoid costly operations (e.g. network updates).
h(o(s)) g) < for some fixed e S + ((A), we can measure the (undiscounted) N-step state and goal occupancies respectively as
2.2. Optimal Transport
N
In the field of machine learning, it is often of interest to match distributions, i.c. find some probability measure that resembles some other probability measure v. In re- cent years there has been an increased interest in Optimal Transportation (OT) (Amos al., 2023; Haldar et al., 2022:
Σ
Pr(s=sd
st]
N +1
or reward-free free Markov Decision Process.
where (SS denotes the set distributions over
2

---

Zero-Shot Offline Imitation Learning via Optimal Transport
in GC-RL. We can introduce this framework by enriching the controllable Markov Chain M. We condition it a goal g E 9 and cast it as an (undiscounted) Markov Decision Process Mg (S ( pg. Poo Rg, Tmax). Compared to the reward-free setting above, the dynamics now include a sink-state I upon goal-reaching and a reward of - until this happens:
Figure 2. An example of Optimal Transport between the discrete approximation f.. & of two Gaussians Je, v. The cost matrix C consists of the pose costs where the cost here is the Euclidian distance. A coupling matrix T' E L( P) (middle) is visualized through lines representing the matching (right).
P(s,a) ifh(o(s) 2e
(5(8,0)
81
otherwise
-11
1 h()(()))))))
Bunne et al., 2023; Pooladian et al., 2024). As illustrated in figure 2, OT does not only compare probability mea- sures in a point-wise fashion, like f-Divergences such as the Kullback-Leible Divergence (DxL), but also incorpo- rates the geometry of the underlying space. This also makes OT robust empirical approximation (sampling) of proba- bility measures ((Peyré & Cuturi, 2019), p.129). Formally, OT describes the coupling E P(X' x )²) of two measures E P E P())) with minimal transportation cost w.r.t. some cost function c : X' x y + IR. The primal Kantorovich form given as the optimization problem
(g a) =
(8)
0
otherwise
where 85 stands for the probability distribution assigning all probability mass to z.
We can now define the goal-conditioned value function as
Piaxx
(T0.
E Ho. Pa
Σ
Rg (12, (e), az)
(9)
where 80 ~ 40: 8t+1
(8 ar); at ~ (((.. The optimal goal-conditioned policy is then T arg max, y Egreg..sppo V/ (*o::) g) for some goal distribution ng E I(G)). Intuitively, the value function VT g) cor- responds to the negative number of expected steps that needs to move from state to goal g. Thus the distance d = V corresponds to the expected first hit time. If no goal abstraction is present, i.e. ó = ids, then (S, d) is a quasimetric space (Wang et al., 2023), i.c. is non-negative and satisfies the triangle inequality. Note, though, that d does not need be be symmetric.
/
V))(U,)= inf vEll(p.)) xxy
22)),
(4)
where the optimization is over all joint distributions of μ and y denoted 7 U(, (couplings). If X' = y and (X,c) is metric space then for E N, WF- = OT is called the Wasserstein-p distance which was shown to be a metric on the subset of measures on X with finite p-th moments (Clement& & Desch, 2008).
Given samples 21,-..,TT ~ f and y1.-..; Vm the dis- crete OT problem between the discrete probability measures >i 1 a;O., and D = 2 bydn, can be written as discrete version of equation 4, namely
3. Goal Abstraction and Myopic Planning
The distribution matching objective at the core of IL prob- lems is in general hard to optimize. For this reason, tra- ditional methods for zero-shot IL leverage a hierarchical decomposition into a sequence of GC-RL problems (Pathak et al., 2018; Hao et al., 2023). We will first describe this ap- proach, and then show how it potentially introduces myopic behavior and suboptimality.
n m
(T, = min b)
ΣΣ (ri,y;)TT
1=15=1
min TEl((a,b)
(C,T)
(6)
with the cost matrix Cıy c(z,, yj) The marginal con- straints can now be written as U(a b) = {TT E Raxm T 1m = a and T 1, = b). This optimization problem can be solved via Linear Programming. Furthermore, Cu- turi (2013) shows that the entropically regularized version, commonly given as OT D) minreu(a..)) (C, T) mDKL (T, ab") can be efficiently solved in its dual form using Sinkhorn's 's algorithm (Sinkhorn & Knopp, 1967)
In the pretraining phase, Pathak et al. (2018) propose to train a goal-conditioned policy Ny SS x g > A on reaching single goals and a goal-recognizer C S that detects whether given state achieves the given goal. Given an expert demonstration 91:M and an initial state So. 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, T. chooses actions to reach the current goal; on the higher level. C' decides whether the current goal is achieved and ny should target the next goal in the sequence.
+ {0,
2.3. Godd-conditionedd Reinforcement Learning
As techniques from the literature will be recurring in this work, we provide short introduction to fundamental ideas
3

---

Zero-Shot Offline Imitation Learning via Optimal Transport
4. Optimal Transport for Zero-Shot IL
(go)
Armed with recent tools value estimation, model-based RL and trajectory optimization, we propose method for zero-shot offline imitation learning that directly optimizes the occupancy matching objective, introducing only mini- mal approximations. As result, the degree of myopia is greatly reduced, as we show empirically in section 5.
start
(0.0)
(1.0)
ag
(gh)
(9a)
1-p
p
(1.1)
(2.1) 1)
In particular, we propose to solve the occupancy matching problem in equation 3 by minimizing the Wasserstein-1 metrie W1 with respect to goal metric h on the goal space G, i.e.
Figure 3. Controllable Markov Chain with @ (a. 1/ + x.
We define the pre-image و (g) = {s S (s)) 9) as the set of all states that map to a goal, and formalize the suboptimality of the above method under goal abstraction as follows.
WI PN) (PN PN).
(11)
Proposition 3.1. Let us define the optimal classifter C(s,g) = 1)) Given a set of visited states P C S, the current state 8 P, and a goal sequence 91:M E GMM, let the optimal hierarchical policy be (s) = (s.gi+1) where i is the smallest integer such that there exist state Sp E P with h((((p)) [K) < and i = 0 otherwise. There exists controllable Markov Chain M and a realizable sequence of goals 9o:M such that, under a suitable goal abstraction o()), Th will not reach all goals in the sequence, 2 (g:) for some E [o] and all N N. i.e. PN
This objective involves two inaccessible quantities: goal occupancies PT: 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 seal- able deep RL techniques.
Occupancy Estimation Since the expert's and the pol- icy's 's occupancy are both inaccessible. opt for discrete, sample-based approximations. In the case of the expert occupancy pN: the single trajectory provided at inference (go...., 9M)) represents a valid sample from it, and we use it directly. For an arbitrary agent policy n, we use a discrete approximation after training a dynamics model p P on Dg, which can be done offline through standard supervised learning. We can then approximate piv by jointly rolling out the learned dynamics model and the policy æ. We thus get the discrete approximations
Proof. Consider the Markov Chain M depicted in figure 3 with goal abstraction & (3) + x and p > O. Now, consider the goal sequence (go. (90.91.92) (1.99119) 92)
(0, 1, which
can only be achieved, by policy taking action 01 in the initial state so = (0, 0). Consider गे in So. with P = { The smallest integer such that ((((5)) 9:) < 6 is i = O, therefore =x())) = π (s0. 91). We can then compare the state-action values Q in sp:
M
F
RJ
PM
8g
and
(12)
M
j=0
T max
-p' 1 - (Tmox (Tmmx+)) +1)
N Sp(ee)
Qr"(-o0) 01,
2
p5y Rd pN
(13)
1 p
N
1=0
<1 = Q7' (1.1) ao 91). (10)
where for the latter we sample
This implies that (50) 7T 1) ap- The next
so po, S141 P(st, as ~ (s)).
(14)
state visited by क will always be 0). from which (2, 1) is not reachable, and J2 is not achievable. We thus have p (92) = 0 for all N e N.
Similarly, we can also obtain an estimate for the state occu- 1 Ei= de.-
pancy of x as ow x 0%
N+1
Metric Approximation As h may be unavailable or hard to specify in practical settings, we propose to train a goal- conditioned value function V from the offline data D3 and use the distance d(s. g) V g) (i.e. the learned first hit time) as a proxy. For a given state-goal pair (s,9), this corresponds to the approximation d(s, g) x h(o(s).) g). It is easy to show that a minimizer of (((()) g) also minimizes d(..g). Using d also has the benefit of incorporating the dynamics of the MDP into the cost of the OT problem. The
We remark that this issue arises in the presence of goal abstraction which plays vital role in the partial demon- stration 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 + 0 as well). Nevertheless, goal abstraction is ubiquitous in practice (Schaul et al., 2015) and necessary to enable learning in complex environments (Andrychowicz al., 2017).
4

---

Zero-Shot Offline Imitation Learning via Optimal Transport
use of this distance has seen some use as the cost function in Wasserstein metrics between state occupancies in the past (Durugkar et al., 2021). 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 / : 9 x g + IR is a distance in goal-space, d : S x g - IR is a distance between states and goals. Nonetheless, d remains applicable as the policy's 's occupancy can also be estimated in state spaces o Given the above considerations, we can rewrite our objective as the discrete optimal transport problem
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 D' ( For i E [o..... MJ. we can then estimate the time when gi should be reached as
o(s")) - (ss; ((s)(s)))) {)
4 R (s00: go)
Σ W(()-1:9))-
(17)
We then simply truncate the online problem to only con- sider goals relevant to $1...... H i.. 90. .o.....99 9K where K = min(j : ty 2 k+ H+ We note that this approximation of the infinite horizon objective can potentially result in my- opic behavior K < M, nonetheless, optimal behavior is recovered as the effective planning horizon increases. Algo- rithm shows how the practical OT objective is computed.
π = arg min OTa(0%
(15)
Optimization Having addressed density and metric ap- proximations, we now focus on optimizing the objective in equation 15. Fortunately, as a discrete OT problem, the objective can be evaluated efficiently using Sinkhorn's 's algo- rithm when introducing entropic regularization with a factor 7 (Cuturi, 2013; Peyré& & Cuturi, 2019). A non- non-Markovian, deterministic policy optimizing the objective at state 8 E S can be written as
Algorithm 1 OT cost computation for ZILOT
Require: Pretrained GC value functions V, W and dynam- ics model P; horizon H, solver iterations and regular- ization factor v)-
Initialization: State so and expert trajectory II:M, precom- puted ta:M
(16)
M
input State history Suck. future actions Hk:
23 arg min min dxi L.N
OTam
80:
N+I
M+I
(Rollout learned dynamics]
==
8k Sk+1ikH + rollout (P, Skk -
{Compute which goals are reachable]
where Susk are the states visited so far and 8k+1:N are rolled out using the learned dynamics model f and actions -I- Note that while 8o:k are part of the objective, they are con- stant and are not actively optimized.
K min(j 1 2 k + H}
(Compute cost matrix)
C₁ + VV 9;); i E {0,.... .....,,.HHH ++H E K}
{Compute uniform marginals}
Intuitively, this optimization problem corresponds to finding the first action from a sequence (at: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., 2018; Ha & Schmidhuber, 2018); while these algorithms are traditionally used for additive costs, the flexibility of zero-order optimizers (Rubinstein & Kroese, 2004; Williams et al., 2015; Pinneri et al., 2020) allows straightforward application to our problem. The objective in equation 17 can thus be directly optimized with CEM variants (Pinneri et al., 2020) or MPPI (Williams et al., 2015), in a model predictive control (MPC) fashion.
a f k+#++ 1 1 +++++1, b + K 1x+1
{Run Sinkhorn Algorithm]
T + siinkhorn(a (a, b, C, r,,,) output TiyCCa {Return OT cost)
The method presented relies solely on three learned components: a dynamics model P, and the state-goal and goal-goal GC value functions V and W. All of them can be learned offline from the dataset Dz. 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., 2024), 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 s encoders. We adapt the method to allow estimation of goal-conditioned value functions, as described in appendix C. We follow prior work (Andrychowicz et al., 2017: Bagatella & Martius, 2023; Tian et al., 2021) and sample goals from the future part of trajectories in D3 in order to synthesize rewards without supervision. We note
Implementation
Like for other MPC approaches, we are forced to plan for finite horizon H, which might be smaller than N, because of imperfections in the learned dynamics model or compu- tational constraints. This is referred to as receding horizon control (Datko, 1969). When the policy rollouts used for computing Φhv are truncated, it is also necessary to trun- cate the goal sequence to exclude any goals that cannot be
5

---

Zero-Shot Offline Imitation Learning via Optimal Transport
that this goal-sampling method also does not require any knowledge of h.
Trajectories
Plan
Render
5. Experiments
This section constitutes an extensive empirical evaluation of ZILOT for zero-shot IL. We first describe our experimental settings, and then present qualitative and quantitative result, as well as ablation study. We consider selection of 30 tasks defined over 5 environments, as summarized below and described in detail in appendix A and C.
Goal Future (dec)
MPC+CIs
ZILOT (ours)
Past (dec) Coupling
fetch (Plappert et al., 2018) is a manipulation suite in which robot arm either pushes (Push), or lifts (Pick&Place) a cube towards a goal. To illustrate the failure cases of my- opic 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 touching the table. As a result, the agent cannot fully con- strain 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".
Figure 4. Example tasks in fetch.s wtw.ssiideelee iaeegeee22 The left two columns show five trajectories across five seeds of the myopic method MPC+CIs and ZILOT (ours). The trajectories are drawn in the zz-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, visualize the OT objective at around three quarters of the episode time. It includes both the past and planned future states, as well their coupling to the goals. Note that planning occurs in the latent state of TD-MPC2, and separately trained decoders used for this visualization.
hal 1all fcheetah (Wawrzyíski, 2009) is classic Mujoco environment where the agent controls 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 tasks involving standing up and hopping on front or back legs, as well as doing flips.
appendix for a discussion of all FB-IL approaches in our rough and partial setting.
Metrics We report two metrics for evaluating planner per- formance. The first one is the minimal encountered (em- pirical) Wasserstein-1 distance under the goal metric h of the agent's trajectory and the given goal sequence. For- mally, given trajectory (8)....,5x) and the goal sequence (o......M)) we define Wmin (80. MM.MMM ) as
pointmaze (Fu et al., 2021) involves maneuvering a point- mass through 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. (2018), 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 Wmin per environment among the values {1, 2, 3, 4,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+CIs), 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., 2020). We further compare ZILOT to ERyB and RERFB, two approaches that combine zero-shot RL. and reward-based IL using the -aakard- (FB) framework (Pirotta et al., 2024). We refer the reader to
min ke (E.......N
W1
Σ
(e)) M+1
1
Σ J=0
M
(18)
1:
We introduce a secondary metric "GoalFraction", which 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 MPC+CIs and ZILOT imitate an expert sliding a cube across the big table of the fetchhs we..ch..iiio side large 2D environment. The myopic baseline struggles to regain control over the cube after moving it towards the second goal, leading to trajecto- ries that leave the manipulation range. In contrast, ZILOTT
6

---

Zero-Shot Offline Imitation Learning via Optimal Transport
fetch_pick_ and_pladeeeeeeaeeeeeaaaaaaaaaaal
fetch_push
fetch_slide__ant 2D
halfcheetah
cettttmaze_meddiuee
2
0.3
0.3
0.3
W Wmin
0.2
0.2
0.2
1
0.1
0.1
0.1
0.0
0.0
0.0
1.0
1.0
1.0
1.0
1.0
GoalFraction
0.5
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
0.0
ERFa
RERFB
Pi+Cls
MPC++Cls
ZILOT (ours)
Figure 5. Performance comparison ZILOTT other methods aggregated over environments. Table reports more detailed results.
plans beyond the second goal. As displayed in the middle part of figure 4, the coupling of the OT problem approxi- mately pairs up each state in the planned trajectory with the appropriate goal, leading to closer imitation of the expert.
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., 2018; Séjourné al., 2019), which turns the hard marginal constraint U (see section 2.2) into a soft constraint. We use this method to address the fact that rough expert trajectory may not necessarily yield a feasible expert occupancy approximation.
5.2. How does ZILOT perform compared to prior methods?
We provide quantitative evaluation of ZILOT with respect to the other planners table 5. For more details we re- fer the reader to appendix A. As ZILOT directly optimizes distribution matching objective, it generally reproduces expert trajectories more closely, achieving lower Wasser- stein distance to its distribution. This is especially evident in environments that are very punishing to myopic plan- ning, 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 overall closer match of the expert trajectory which is most clearly visible in the poi na environment. We note that the performance of the two myopic baselines Pi+CIs and MPC+CIs are very similar, suggesting that the performance gap to ZILOT stems from the change in objective, rather than implementation or model-based components. We suspect the origins of the sub- par performance of ER and RERFB are two-fold. First, the FB framework (Touati & Ollivier, 2021; Pirotta et al., 2024) has been found to underperform in low-data regimes(Jeen et al., 2024). Second, ERTy and REREB use a regularized f-divergence objective, which translates to an RL problem with additive rewards. As (Pirotta et al., 2024) state, this regularization comes at cost, particularly if states do not contain dynamical information or in ergodic MDPs. In this case, policy can optimize the reward by remaining in the most likely expert state, yielding a degenerate solution. Con- versely, such a solution would be discarded by ZILOT as it uses an unregularized objective.
OT+CIs, version of our method which uses the clas- sifier (Cls) (with the same hyperparameter search) to discards all goals that are recognized as reached. This allows this method to only consider future goals and states in the OT objective.
OT+h, our method with the goal metric h on G the cost function in the OT problem, replacing d.
Our results are summarized in figure 6. First, we see that us- ing 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 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 en- vironments, 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 E. Fur- thermore, while is locally accurate, it ignores the global geometry of MDPs, as shown by its poor performance in strongly asymmetric environments (i.e., hal fcheetah).

---

Zero-Shot Offline Imitation Learning via Optimal Transport
fetch_pick_ and _
fetch_push
@etcchsside__ite sariide__rrg___ee 2D
halfcheetah
bointmazze_ediiie
2
0.3
0.3
0.3
Wmin
0.2
0.2
0.2
1
0.1
0.1
0.1
0.0
0.0
0.0
1.0
1.0
1.0
1.0
1.0
GoalFraction
0.5
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
0.0
ZILOT+h
ZILOT+Cls
ZILOT+Unbalanced
ZILOT (ours)
Figure 6. Ablation of design choices ZII.OT, including coupling constraints (OT+unbalancee), partial trajectory matching (OT+CIs), and the approximation of A by d (OT+h), For more detailed results, please refer to table 2.
6. Related Work
(Ho & Ermon, 2016; Li et al., 2023) or in one instance OT (Luo et al., 2023). Another line of work has shown impres- sive real-world results by matching the action distributions (Shafiullah al., 2022; Florence et al., 2021; Chi et al., 2023) directly. All these approaches do not operate in zero-shot fashion, or need ad-hoc data collection.
Zero-shot II, When a substantial amount of compute is allowed at inference time, several existing methods leverage pretrained models to infer actions, and retrieve an imitator policy via behavior cloning (Pan et al., 2020; Zhang al., 2023; Torabi al., 2018). 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., 2018; Hao et al., 2023). Expert demonstrations are then imitated by going from one goal to the next using a coal-conditioned policy. In contrast, our work proposes holistic approach imitation, which considers all goals within the planning horizon.
OT in RL Previous works have often used OT as a re- ward signal RL. One application is online fine-tuning, where a policy's 's rollouts are rewarded proportion to how closely they match expert trajectories (Dadashi et al., 2021; Haldar et al., 2022). Luo et al. (2023) instead use a similar trajectory matching strategy to recover reward labels for unlabelled mixed-quality offline datasets. Most of the works mentioned above rely on simple Cosine similarities and Euclidean distances as cost-functions their OT problems..
Zero-Shot RL. Vast amounts of effort have been dedicated to learning generalist agents without supervision, both on the theoretical (Touati & Ollivier, 2021; Touati et al., 2023) and practical side (Laskin et al., 2021; Mendonca et al., 2021). Among others, (Sancaktar et al., 2022; P. et al., 2021; Bagatella & Martius, 2023) learn dynamics model through curious exploration and show how it can be leveraged to optimize additive objectives. More recently, Frans et al. (2024) use Functional Reward Encodings to encode arbi- trary 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., 2024) (i.e., distribution matching). One notable exception is the forward-backward framework (Touati & Ollivier, 2021; Pirotta et al., 2024), which we discuss in detail in appendix B.
7. Discussion
In this work, we point out a failure-mode of current zero- shot IL methods that cast imitating an expert demonstra- tion as following 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 Optimal Transporta- tion 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-ther thetthe zero-shot IL methods, particularly scenarios where non-myopic planning is cru- cial. We additionally validate our design choices through a series of ablations.
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 match- ing objectives (Liu et al., 2023; Ma et al., 2022) recover reward using Generative Adversarial Networks (GAN)
From practical standpoint, our method is mainly limited in its reliance on a world model. As the inaccuracy and com- putational cost of learned dynamics models increase with the prediction horizon, we are forced to optimize a fixed- horizon objective. This may reintroduce a slight degree of
8

---

Zero-Shot Offline Imitation Learning via Optimal Transport
myopia that could lead actions which cause suboptimal behavior beyond the planning horizon. This, however, was not practical issue in our empirical validation, and we expect our framework to further benefit as the accuracy of learned world models improves. From a theoretical stand- point, ZILOT induces a non-Markovian policy, even when expert trajectories are collected by a Markovian policy, and a Markovian policy would thus be sufficient for imitation. While the space of Mon-Markovian policies is larger, we find ZILOT to be able to efficiently find near-optimal policy. This aligns with the fact that several practical zero-shot II. algorithms are based on efficient search over non-Markovian policies (e.g. those based on a goal-reaching policy and a classifier (Pathak et al., 2019; Pirotta et al., 2024)).
5048-5058, 2017.
Bagatella, M. and Martius, G. Goal-conditioned offline planning from curious exploration. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans,
LA, USA, December 10 16, 2023, 2023.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020.
Impact Statement
Advancements in imitation learning may lead to more capa- ble robotic systems across variety of application domains. While such systems could have societal implications de- pending on their use cases, our contributions are algorithmic rather than comain-specific.
Bunne, C., Stark, S. G., Gut, G., del Castillo, J. S., Levesque, M.. Lehmann, K.-V., Pelkmans, L.,, Krause, A., and Rätsch, G. Learning single-cell perturbation responses using neural optimal transport. Nature Methods, 20 1611759-176,, Nov 2023. ISSN 1548-7105. doi: 10. 10.10388/441592-23--096666
Acknowledgments
We thank Anselm Paulus, Mikel Zhobro, and Núria Ar- mengol Urpi 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: 011S 18131))))) Finally, this work was sup- ported by the ERC- - 101045454 REAL-RL.
Chi, C., Feng. S., Du, Y., Xu, Z., Cousineau, E.,, Burchfiel, B., and Song. S. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Chua, K., Calandra, R., McAllister, R., and Levine, S. Deep reinforcement learning in a handful of trials using proba- bilistic dynamics models. In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K.,
N., and Gar- nett, R. (eds.), Advances in Neural Information Process- ing Systems 31: Annual Conference on Neural Informa- tion Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 4759-4770, 2018
References
Amos, B., Luise, G., Cohen, S., and Redko, I. Meta op- timal transport. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Inter- national Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 791-813. PMLR, 2023.
Clement, P. and Desch, W. An elementary proof of the tri- angle incquality for the wasserstein metric. Proceedings of The American Mathematical Society- PROC AMER MATH SOC, 136:333-340 01 2008. doi: 10.1090//0002-
9939-07-09020-X
Andrychowicz. M., Crow, D., Ray, A., Schneider, J.. Fong, R., Welinder, P., McGrew, B., Tobin, J., Abbeel, P., and Zaremba, W. Hindsight experience replay. In Guyon, I,, von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp-
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. In Burges, C., Bottou, I.., Welling, M., Ghahramani, Z., and Weinberger, K. (eds.), Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013.
Cuturi, M., Penng-aapxxanthss L., Tian, Y., Bunne, C., Davis, G., and Teboul, O. Optimal transport tools (ott):

---

Zero-Shot Offline Imitation Learning via Optimal Transport
A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201. 123244 2022.
encodings. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net 2024.
Dadashi, R., Hussenot, L.. Geist, M., and Pietquin, O. Pri- mal wasserstein imitation learning. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4RL: Datasets for Deep Data-Driven Reinforcement Learning, 2021.
Ghosh, D., Walke, H., Pertsch, K., Black, K., Mees, O., Dasari, S., Hejna, J., Kreiman, T., Xu, C., Luo, J., Tan, Y. L., Chen, L. Y.. Sanketi, P., Vuong, Q., Xiao, T., Sadigh, D., Finn, C., and Levine, S. Octo: An open-source gen- eralist robot policy. CoRR, abs/2405. 2024, doi: 10.48550/ARXI.2205..00
Datko, R. Foundations of optimal control theory (e. bruce lee and lawrence markus). SIAM Rev. 11(1):93-99) 95, Jan- uary 1969. ISSN 1045-1445. doi: 10. 1 1111/1111100 1020.
Duan, Y., Andrychowicz, M., Stadie, B. C., Ho, J., Schnei- der, J., Sutskever, L., Abbeel, P., and Zaremba, W. One- shot imitation learning. In Guyon, I., von Luxburg. U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.). Advances in Neural In- formation Processing Systems 30: Annual Conference Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 1087-1098, 2017.
Ha, D. and Schmidhuber, J. World models. CoRR, 2018.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor- critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Dy, J. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Ma- chine Learning Research, pp. 1861-1870. PMLR, 10-15 Jul 2018
Durugkar, I, Tec, M., Niekum, S., and Stone, P. Adversarial intrinsic motivation for reinforcement learning. In Ran- zato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, De- cember 14, 2021, virtual, pp. 8622 2636, 2021.
Haldar, S., Mathur, V., Yarats, D., and Pinto, L. Watch and match: Supercharging imitation with regularized optimal transport. In Liu, K., Kulic, D., and Ichnowski, J. (eds.), Conference on Robot Learning, CoRL 2022, 14-18 De- cember 2022, Auckland, New Zealand, volume 205 of Proceedings of Machine Learning Research, pp. 32-43. PMLR, 2022.
Eysenbach, B., Zhang, T., Levine, S., and Salakhutdinov, R. Contrastive learning as goal-conditioned reinforce- ment learning. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Hansen, N., Su, H., and Wang, X. TD-MPC2: scalable, ro- bust world models for continuous control. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenRe- view.net, 2024.
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Bois- bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras,
K., Fournier, N., Gautheron, I.., Gayraud, N. T., Janati, H., Rakotomamonjy, A., Redko, I., Rolet, A., Schutz, A., Seguy, V., Sutherland, D. Tavenard, R., Tong, A., and Vayer, T. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78):1-8, 2021.
Hao, P., Lu, T., Cui, S., Jei, Cai, Y., and Wang, S. Sozil: Self-optimal zero-shot imitation learning. IEEE Trans- actions on Cognitive and Developmental Systems, 15(4); 2077-2088. 2023. doi: 10. 10.1109/TCDS
Ho, J. and Ermon, S. Generative adversarial imitation learn- ing. In Lee, D., Sugiyama, M., Luxburg. U., Guyon, I., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
Florence, P., Lynch, C., Zeng, A., Ramirez, O. A., Wahid, A., Downs, L., Wong, A., Lee, J., Mordatch, I., and Tompson, Implicit behavioral cloning. In Faust, A., Hsu, D., and Neumann, G. (eds.), Conference on Robot Learning, 8-11 November 2021, London, UK, volume 164 of Proceedings of Machine Learning Research, pp. 158 168. PMLR, 2021.
Jang, E., Irpan, A., Khansari, M., Kappler, D., Ebert, F, Lynch, C., Levine, S., and Finn, C. BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning. In Proceedings of the Sth Conference on Robot Learning, pp. 991-1002, PMLR, 2022.
Frans, K., Park, S., Abbecl, P., and Levine, S. Unsupervised zero-shot reinforcement learning via functional reward
10

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Jeen, S., Bewley, T., and Cullen, J. M. Zero-shot reinforce- ment learning from low quality data. Advances in Neural Information Processing Systems 38, 2024.
Ma, Y. J., Shen, A., Jayaraman, D., and Bastani, O. Versa- tile offline imitation from observations and examples via regularized state-occupancy matching. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvári, C., Niu, G., and Sabato, S. (eds.). International Conference on Machine Learn- ing, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 14639-14663 PMLR, 2022.
Kim, G., Seo, S., Lee, J., Jeon, W., Hwang. H., Yang, H., and Kim, K. Demodice: Offline imitation learning with supplementary imperfect demonstrations. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenRe- view.net, 2022.
Mendonca, R., Rybkin, O.,, Daniilidis, K., Hafner, D., and Pathak, D. Discovering and achieving goals via world models. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neu- ral Information Processing Systems 34: Annual Confer- ence on Neural Information Processing Systems 2021, NewrIPS 2021, December 6- 14, 2021, virtual, pp. 24379- 24391, 2021.
Kim, M. J., Pertsch, K., Karamcheti, S., Xiao, T., Bal- akrishna, A., Nair. S., Rafailov, R., Foster, E. P.. Lam, G., Sanketi, P., Vuong, Q., Kollar, T., Burchfiel, B., Tedrake, R., Sadigh, D., Levine, S., Liang, P. and Finn, C. Openvla: An open-source vision-language- action model. CoRR, abs/2406.09246. 2024. doi: 10.4850/ARXIV.446600222
O'Neill, A., Rehman, A., and et. al. Open X-Embodiment: Robotic learning datasets and RT-X models. https: / //arxiv.org.e hww..oor//bbao... /abg///b////// /aa.////1000 231 0. 088 2023.
Kojima, T., Gu, S. S., Reid, M., Matsuo, Y., and Iwasawa, Y. Large language models are zero-shot reasoners. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances Neural Informa- tion Processing Systems, volume 35, pp. 22199-22213. Curran Associates, Inc., 2022.
P., M. V., Blaes, S., Pinneri, C., and Martius, G. Risk- averse zero-order trajectory optimization. In Faust, A., Hsu, D., and Neumann, G. (eds.), Conference on Robot Learning, 8-11 November 2021, London, UK, volume 164 of Proceedings of Machine Learning Research, pp. 444 PMLR, 2021.
Laskin, M., Yarats, D., Liu, H., Lee, K., Zhan, A., Lu, K., Cang, C., Pinto, I.., and Abbeel, P. URLB: unsupervised reinforcement learning benchmark. In Vanschoren, and Yeung, S. (eds.), Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1, NeurlPS Datasets and Benchmarks 1021, December 2021, virtual, 2021.
Pan, X., Zhang, T., Ichter, B., Faust, A., Tan, J., and Ha, S. Zero-shot imitation learning from demonstrations for legged robot visual navigation. In 2020 IEEE Interna- tional Conference on Robolics and Automation, ICRA 2020, Paris, France, May 31 August 31, 2020, pp. 679-685. IEEE, 2020. doi: 10. 1011009////444445.20000 2020. 9196602.
Li, C., Vlastelica, M., Blaes, S., Frey, J., Grimminger, F., and Martius, G. Learning agile skills via adversarial imitation of rough partial demonstrations. In Conference on Robot Learning, pp. 342-352. PMLR, 2023.
Pathak, D., Mahmoudieh, P., Luo, G., Agrawal, P., Chen, D., Shentu, Y., Shelhamer, E., Malik, J., Efros, A. A., and Darrell, T. Zero-shot visual imitation. In 2018 IEEE Conference on Computer Vision and Pattern Recognition Workshops, CVPR Workshops 2018, Salt Lake City, UT, USA, June 18-22, 2018, pp. 2050-2053, Computer Vision Foundation IEEE Computer Society, 2018. doi: 10. L109/CVPR.22018.000 8.00278..
Liero, M., Mielke, A., and Savaré, G. Optimal entropy- transport problems and new hellinger-kantoroviie dis- tance between positive measures. Inventiones mathemati- cae, 211, 03 2018. doi: 10. 10.1007//00222-007--07999
Liu, J., He, L., Kang, Y., Zhuang, Z., Wang, D., and Xu, H. CEIL: generalized contextual imitation learning. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, IA, USA, December 10 16, 2023, 2023.
Pathak, D., Gandhi, D., and Gupta, A. Self-supervised exploration via disagreement. In ICML, 2019.
Peyré, G. and Cuturi, M. Computational optimal transport. Found. Trends Mach. Learn., 11(5-6)3355-677, 2019, doi: 10.1561/2200000073..
I Luo, Y., Jiang, Z., Cohen, S., Grefenstette, E.,, and Deisen- roth, M. P. Optimal transport for offline imitation learning. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
Pineda, L., Amos, B., Zhang, A., Lambert, N. O., and Ca- landra, R. Mbrl-lib: A modular library for model-based reinforcement learning. Arxiv, 2021.

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Pinneri, C., Sawant, S., Blaes, S., Achterhold, J., Stueck- ler, J., Rolínek, M., and Martius, G. Sample-efficient cross-entropy method for real-time planning. In Kober, J.. Ramos, F., and Tomlin, C. J. (eds.), 4th Conference on Robot Learning, CoRL. 2020, 16-18 November 2020, Virtual Event Cambridge, MA, USA, volume 155 of Pro- ceedings of Muchine Leurning Research, pp. 1049-1065.
submodular semi-gradient methods. In Forty-first Inter- national Conference on Machine Learning, ICML. 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024.
Schaul, T., Horgan, D., Gregor, K., and Silver, D. Universal value function approximators. In Bach, F. R. and Blei, D. M. (eds.), Proceedings of the 32nd International Con- ference on Machine Learning, 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Con- ference Proceedings, pp. 1312-1320. JMLR.org. 2015.
PMLR. 2020.
Pirotta, M., Tirinzoni, A., Touati, A., Lazaric, A., and O1- livier, Y. Fast imitation via behavior foundation models. In The Twelfth International Conference on Leurning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
Séjourné, T., Feydy, J., Vialard, F., Trouvé, A., and Peyré, G. Sinkhorn divergences for unbalanced optimal transport. CoRR, abs/19110 12958, 2019.
Shafiullah, N. M., Cui, Z. J., Altanzaya, A., and Pinto, L. Be- havior transformers: Cloning Sk$ modes with one stone. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Infor- mation Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Plappert, M., Andrychowicz, M., Ray, A., McGrew, B., Baker, B., Powell, G., Schneider, J., Tobin, J., Chociej, M.. Welinder, P., Kumar, V., and Zaremba, W. Multi-goal reinforcement learning: Challenging roboties environ- ments and request for research, 2018.
Pooladian, A., Domingo-Enriehe Enrich, C., Chen, R. T. Q.. and Amos, B. Neural optimal transport with lagrangian costs. CoRR, abs/2406.00028,, 2024. doi: 10.48550/ARXIV. 2406.00288.
Sinkhorn, R. and Knopp, P. Concerning nonnegative ma- trices and doubly stochastic matrices. Pacific Journal of Mathematics, 21:343 348, 1967.
Reddy, S., Dragan, A. D., and Levine, S. SQIL: imitation learning via reinforcement learning with sparse rewards. In 8th International Conference on Learning Represen- tations, ICLR 2020, Addis Ababa, Ethiopia. April 26-30, 2020. OeenRevi, 2020.
Sutton, R. S. and Barto, A. G. Reinforcement Learning: An Introduction. A Bradford Book, Cambridge, MA, USA, 2018. ISBN 0262039249.
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., de Las Casas, D., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., Lillicrap. T., and Riedmiller, M. Deepmind control suite, 2018.
Reed, S. E., Zolna, K., Parisotto, E., Colmenarejo, S. G., Novikov, A., Barth-Maron, G., Gimenez, M., Sulsky, Y., Kay, J., Springenberg. J. T., Fecles, T., Bruce, J., Razavi, A., Edwards, A., Heess, N., Chen, Y., Hadsell, R., Vinyals, O., Bordbar, M., and de Freitas, N. A generalist agent. Trans. Mach. Learn. Res., 2022, 2022.
Tian, S., Nair, S., Ebert, F., Dasari, S., Eysenbach, B., Finn, C., and Levine, S. Model-based visual planning with self-supervised functional distances. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Rubinstein, R. Y. and Kroese, D. P. The Cross Entropy Method: A Unified Approach To Combinatorial Opti- mization, Monte-carlo Simulation (Information Science and Statistics). Springer-Verlag, 2004. ISBN 978-0-387- 21240-1.
Torabi, F., Warnell, G., and Stone, P. Behavioral cloning from observation. In Lang, J. (ed.), Proceedings of the Twenth-Sevetth International Joint Conference on Arti- ficial Intelligence, IJCAI 2018, July 13-19, 2018, Stock- holm, Sweden, pp. 4957. ijcai.org, 2018. doi: 10. 00224666/ICCAA220188688
Sancaktar, C., Blaes, S., and Martius, G. Curious explo- ration via structured world models yields zero-shot object manipulation. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information P Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022, 2022.
Touati, A. and Ollivier, Y. Learning one representation to optimize all rewards. In Ranzato. M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6- 64, 2021, vir- tual, pp. 13-23, 2021.
Santi, R. D., Prajapat, M., and Krause, A. Global reinforce- ment learning: Beyond linear and convex rewards via
12

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Touati, A., Rapin, J., and Ollivier, Y. Does zero-shot rein- forcement learning exist? In The Eleventh International Conference Learning Representations, ICLR 2023, Ki- gali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
Wang, T., Torralba, A., Isola, P., and Zhang, A. Opti- mal goal-reaching reinforcement learning via quasimetric learning, 2023.
Wawrzyúski, A cat-like robot real-time learning to run. In Kolehmainen, M., Toivanen, P., and Beliczynski, B. (eds.), Adaptive and Natural Computing Algorithms, pp. 380-390, Berlin, Heidelberg. 2009. Springer Berlin Hei- delberg. ISBN 978-3-642-04921-1
Wei, J., Bosma, M., Zhao, V. Y., Guu, K., Yu, W., Lester, B., Du, N., Dai, A. M., and Le, Q. V. Finetuned language models zero-shot leamners In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022.
Williams, G., Aldrich, A., and Theodorou, E. A. Model predictive path integral control using covariance variable importance sampling. CoRR, abs/1509.00 149, 2015.
Yang. R., Lu, Y., Li, W., Sun, H., Fang. M., Du. Y., Li, X., Han, L., and Zhang, Rethinking conditondt supervised learning and its connection to offline RL. In International Conference on Learning Representations, 2022.
Zhang, X., Bocker-Ehmeke P., van der Smagt, P., and Karl, M. Action inference by maximising evidence: Zero-shot imitation from observation with world models. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10- 16, 2023, 2023.
Zolna, K., Novikov, A., Konyushkova, K., Gülçehre, Ç.,, Wang, Z., Aytar, Y., Denil, M., de Freitas, N., and Reed, S. E. Offline learning from demonstrations and unlabeled experience. CoRR, abs/201 1. 13885, 2020.
13

---

Zero-Shot Offline Imitation Learning via Optimal Transport
A. Additional Results
A.1. Main Result Details
In tables I and we provide detailed results for the figures 5 and 6. We also provide a summary of all planners we evaluated in figure 7.
letch_pick_ant and place
fetch_push
setch_sside.. lare_ _2
halfcheetah
pointtmaze_meiiiee
2
0.3
0.3
0.3
Wiin
0.2
0.2
0.2
0.1
0.1
0.1
0.0
0.0
0.0
1.0
1.0
1.0
1.0
1.0
GoalFraction
0.5
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
0.0
ERFB
Pi+Cls MPC+CIs
ZILOT+h ZILOT+CIs
ZILOT Unnbalandd
RERyB
ZILOT (ours)
Figure 7. Summarized performance of all discussed Planners. See table and table 2 for detailed results.
Table I. Performance of Pi+Cls, MPC+CIs 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 metries across 5 seeds. We perform Welch f-test with p = 0.05 do distinguish the hest values and mark them hold. Values are rounded to 3 and 2 digits respectively.
Wain 4
GaulFraction PicCls
ERr
PacCs
MPC-CIs
ZILOT (ours)
ERya
RBKr
MPC+CIs
ZILDT (ours)
Fetchipi ol place & denae
0.233=00017
0:233=0::11
0.089+0.027
0.109=0024
0.049:00.019
0.17=0.03
0.17+0.03
0.65+-0.11
0.58+00007
0.88=0.07
ratch panddp.cco----ppre... feichhgi ok pkandd_plaee 8 dense
0.170::0011
0.100=00011
0.112:0:114
0.127=0002
0.09210.015
0.350004
0.3510004
0.6210005
0.65:0.05
0.185+00019
0.183-0.011
0.113+0:022
0.101=0022
0.049:000114
0.16-0004
0.16+00.4
0.1+00.07
0.62+0.08
0.55-0.08
tet
0.09K:000088
0.088000088
0.81100.017
0.091:00007
0.067:0.006
0.33-008
0.331:0.88
0.57:00.66
0,5010.04
0.70::0.06
fatch pick and placo-U-dense
0.121110021
0.12410021
0.12710.007
0.11610015
0.06810.005
0.2110008
0:4710.10
0.6010.03
0.7010.02
teichipi wwc.pppccknddd.... plaand_ppace DD aparss
0.166=00008
0.163=0.028
0.142:0:005
0.160=00000
0.098:000003
0.30=0.07
0.30=+0.07
0.51+00.02
0. 0.38+00.0
0.55=00.05
0.162.10010
0.162100010
0.11110.007
0.117L00022
0.07010.009
0.25_0004
0.25100.44
0.54.10.02
0,52100.2
0.72.0.04
fscoh.push---denee
0,24610:001
0.16600.011
0.056.10.00
0.085L00018
0.01110.015
0.15_0000
0.1510.00
0,9610.03
0.72.10.09
091100.6
fezch, push---spprse
0.184:20:04
0.184=0.144
0,101:000011
0,10:=00100
0.082:1.004
0.35:004
0.35
0.65100009
0.44::0.04
0.69:00066
fctch.puth-S-ddnee
0,182100019
0.182.0019
0.07710.024
0.10110.026
0,04910.010
0.2500055
0.2510.05
0.8310009
0,7010.05
0.57.0.06
Bucch.push---sppae
0.123=00010
0.123=00000
0.062:0.004
0.077=0004
0.06410.006
0.36=0088
0.36:0008
0.90=0.07
0.72=0006
feiolugush-U-deene
0,141+0001
0.141-0011
0.102-0.044
0,001-00009
0.065+0.004
0.28-0.07
0.28+0.07
0,72+0.18
0.67+0.08
6.77-0.02
tetoh.puth----sprre
0.005
0.095.0022
0.10610.014
0.131100112
0,10910.007
0.32 0.03
0.3210.03
0,7011012
0,4510055
0.53.00.3
teichupush a11
0.178=00008
0
0.084:0.007
0.08=0010
0.068:00005
0.29-0.02
0.20
0.7920.05
0.61=000
0.75=0.03
feichuslidee 1. denae
0.282=0014
0.202-0014
0.25==0.022
0,217=0034
0.074+0.011
0.17=0.04
0.17=0.04
0.26++0.06
0.40+0.111
0.76=0003
0:255=00007
0:255000077
0.233-00014
0.185=00027
0.120:0.011
0.38-005
0
0.47.0.10
0.70:0055
0.73:-004
0,232110029
0.29910.006
0.2541.0022
0.19.0055
0.1910005
0.1110.10
0.5110007
tetch_slidee.lorge2222 23-8-ppreeeeee
0211100014
0015:00114
0.266:0.006
0.250
0.06610.015
0.28-0004
0.28:0004
0.3110.02
0:43:00.02
6.7400.M
fatch siide large 23-U-dense
0,22510051
0226000051
0.21410.029
0.191100445
0.076.10.009
0.14-003
0,1410.03
0.3010.07
0,3510.10
0.76.0.44
Geichhallide_14sep222 U. Dpaparse
0.21=00.09
0.204:00048
0:169:0:443
0.150=0012
0.120=40.005
0.30=004
0.20=0004
0.36+0009
0.53+0004
0.7000.06
feichualideelasge.. al:
0.20000:112
0:251=0.012
0:288=0.008
0,205=00020
0.098:000077
0.24-0.02
0.24:0.02
0.32+0.04
0.45+0.04
6.70=0.02
balfchettat-bcckflip
2..02.0.149
2.022...14
1.089
4.281L03371
100750
0.44
0.44
0.28.10.13
0.1210.12
0.57.0.17
2853=20.04
2:83=00104
2879100427
3.04420.722
2.17110.454
0.06::0055
0.4410.10
0.11
Balfcheetah Front: Fllpp
1.286=0.69
1.286-0.059
1.544+0.127
1.655-01177
1.205+0.094
0.73-0:22
0.73+0022
0.77+0009
B. 0.79+0.12
1.00-0.00
2.1
0.13
2.86220133
2,083=01100
1.555
027=006
0.27:0.08
0.70+0.088
0.1000077
0.83-0003
0.910+0.16
0910-00316
0.006+0.110
0.950-0.075
0.509-0.107
0.65+0.19
0.96+-003
0.90+0.02
0.96-001
ba boltcaeetatthoo-teee -ooetatt-hop-tewardd
1418 0.332
0.332
1.50100.099
1,392 0.206
1110
0.13_000
0.4310009
0.5110.07
0.6211.144
0.12
0.607:00079
0.65700079
0.89710.092
0.679=00035
0.46910.167
0.110009
0.8120009
0.9610.04
1.00
0.99:0.01
hal fcheetah run forward
0.712+00044
0.712-0.064
0.857+0.044
0.822-0
0.376+200199
0.76=0.07
0.76+0007
1.0+0.011
0.94:1088
1.00-0.00
halfchectat-oll
1.98110.105
1.18800.105
L71710.101
1.869.000099
1.32510.123
0.50_00.03
0.5210.03
0.7010.00
0.71100.02
0.02
peintmazo. wnw.indddddddd crd.un-cirrle-eeeee
252
1128_00220
0,21310.0038
0.221
0,15610.0010
0.18_00066
0,1810.06
1.0010.00
1.0010.000
1.000.000
pointmaze mediut circle- apazau
1483=0010
1483=0440
0.385=0.015
0.404:0.025
0.466=20.024
0.22=0.00
0:2250.00
1.000..00
1.00:0.00
0.1=0011
0.90010.317
0.10
0.27510.003
0,23810023
0.19910.013
0.56
0.5610.11
1.0010.00
1.0010.000
1.0010.00
pointnaze_nediur sparse
1.0%6=0.505
1:066=0:055
0.555=0.0000
0.511=00055
0.459:0.015
0.48:0.10
0.48=0.10
1.00+0.000
1.00=00.00
0.97=0.03
pointmaze all
1.149#0.163
1.149=0.163
0:36540:021
0.543=00023
0.320=20.009
0.36=00.2
0.36+0.02
1.0000000
1.00:40.00
0:4-00004
14

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Table 2.
Performance of our method and its ablations in all environments and tasks. Fach 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 3 and 2 digits respectively. Task
Wain
GoalFraction 1
ZILOT++
ZILOT=CB
ZILLOT+Unkalancca
ZILOT (ours)
ZILOT+A
ZILOTHCIs
ZIIUTT+Unsslaneed
ZILOT (ours)
0.214
0.09110.011
0.052.0.018
0,019-0,019
0.2610.10
0.6810.04
0,8410.07
0.88
0.18
0.158:0.004
0.09610.0016
0.092=0.015
0.40:0001
0.35::0.02
0.65::0.08
0.65::0.05
0,198-0.42
0.089=0001
+0,06
0.049-0.014
0.36+0 15
0.71+00.77
0..6-+0.03
0185--0.18
settch.iick.and www.dn..ace---ppee.....
0.17:=0.229
0.115200009
0.056:10.008
0.067=0.006
0:42:10.08
0.57:10.02
0.76:1:0008
0.70:00.06
fet ceecchiee send sac.tppplce------- ace--neens
0.237-0.043
0.071+00006
0.060:00.00
(.068 0.005
0.17+0 10
0.74:0004
0.75++0.04
0.70==0.02
fetch.pick
0.229=0.034
0.167:0.004
0.101:00.008
0.098
0.34::0.01
0.33::0.05
0.54::0.05
0.55:20.05
0.207=00026
0.11520.007
0.068:00.008
0.090=00009
0.32::0.06
0,56::0.02
0.73:1005
0.72::0.14
fetch.push-I-denee
0.21=00.200
000
0.040:000044
0.011=0.015
0.27::0.06
0.73:0002
0.91:00.3
0.91::0.06
fetch.pust-i-sparee
0.200
0.15010.005
000011
0,082_00004
0,39100.06
0,36100.3
0.6510..77
0.69
fetch.push----een
0.203=0.0446
0.077:00008
0.009:0.0000
0.049=0.010
0.32:00.14 14
0.72::0.05
0.86::0.05
0.87::0.06
fetchhgush-S-spirre
0,197-0.055
0.097-00066
0.060+00.009
0,064 0.,00
0.40+00117
0.56+00.22
0.78 10.06
41.26
fetcch.uushU--denee
0.228=0.045
0.068:00007
0.00%:00.009
0.065-0.004
0.20::0.10 10
0.78:00.04
0.81:
0.77::0.02
fetcch.push---spareee
0,224-0.847
0.156+0.017
0.100+00.007
0.109-0.007
0.36+0.07
0,39+0.05
IJ61++0.05
0.53-+0.03
fetch.push-all
0.211=0.033
0.10010.0066
0.068.10.005
0.068_00005
0.32:00.88
0.591:0.02
0.77.10.03
0.75
fetch.slide.large
(ataargee0------ 20-2-dense
0.255=0.022
0.09810.0227
0.060:10.009
0.071=00011
0.26110.08
0.6910008
0.81
0.76:10.03
0.236=0.022 0200
0.18
0.112:00.016
0.120=0.011
0.41:0.04
0.45=0008
0.83::0.08
0.73==0.04
fetch.elide..arge.20
0.256-0.035
0.10510.011
0.09110.0099
0.111=0.010
0.2310.100
0.63:0.03
0,59110.10
0.51:10.00
0.272:000445
0.152+0033
0.,00
0.01550.055
0.28:00077
0.52:00088
0.79::0.04
0.74=00.04
0,315-0.851
0.057+00.09
0.074+00001
0.076- 0,009
0.12++0.8
0.75+00007
0.75++0.04
0.76-40.84
0.288=0.058
0.147:00009
0.117::0.008
0.120=0.005
0. 0.30:0004
0.41:0.04
0.68:
0.70::0.06
fetchh.llide..arge.20---
0.270=0.025
0.125:00011
0.0905
0.09
0.27:00004
0,57=00.44
0,74+1.02
O. 0.70==0.02
hal fchheeththh-aackf11
0.317
3,170+0730
2.710+0.744 742
0.780
0.50:0.18
0.43=0.14
0.55
0.57
hal wwc.chetabbbbcckll------eee
2.537=0.810
2.47910.284
10,57
2.171=0,454
0.47:0..77
0,50:0.11
0,56:10.16
0.58:0.11
hal fcheetah-front flipp
1.172-0.091
1.796+0.173
10.10
0.295
0.96::0.03
0.52+0.83
0.9%
1.00:0.00
hal Srreettah--root.llppppp
2.526=0.110
2.091:0.210
1.96910.075
1.955-0.057
0.13:0007
0.60::0.06
0.88:
0.85
hal fcheeteh-hop-backeee
0.739-0.736
0.889+0.103
0.548 0.055
0.107
0.84:0.33
0.22+0007
0.96:10.04
0.96
hal wcheetahh-oo--rrreee
0.12
1.070+0.006
00004
1.101=0.152
0.78:00.12
0.63++0.98
0.67
0.58-=0.12
ha1 fcheetah-run-backwer
0.455-0.415
0.538+0.139
0.473+00162
0.199
0.92:00.11
0.68+00.03
0.99:00.01
0.99:00.1
hal Saneeetth-run-oooeee
0.172-0.156
0.742#0.044
0.381 00006
0.019
0.93:00.09
0,72+0005
1.00+0.001
1.00:00.00
hal Echeetah-ell
0.116:0.181
1.64420.899
1.399:0.000
0.125=0.123
0.691:0.06
0.61:00022
0.83::0.02
0.82:00002
po rointmaze.red un-ci ci rele-dense
0,252=0.832
0.61::0.777
0.165:000115
0.156=0.010
0.91:0004
0.62::0.25
1.00:0.00
1.00::0.00
po mntuendd un-ol é-epa
0.465-0.056
1.074+0.115
0.465 +0.02
0.466 0.024
0
0.41+0.10
0,83++0.10
0.81 0.81-0.111
pointmaze.ned unn pat h-dense
0,495=0.130
1.835:1.064
0.192::0.008
0.199-0.013
0.95:0003
0.45:0.29
1.00:00000
1.00::0.00
po nntaaneeeedd un-palt tem-path-epaaee reeparee
0.716-0.119
1.41++0828
0.444+0.0010
0.459 0.015
0.89+0.10 10
0.61+0.24
0.99+00.0
0.97-00.03
po poi'nn inmameeeeed un-a11
0.482-0.055
1,244+0143
0.317+00008
0.320 0.0009
0.91+00.2
0,52+00115
0.95++0.03
0.94--0.14
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 dynamies 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 8 for environment fetch s diide_laal lar_larre_220 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.
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 9 shows the success rate of all planners in this task verifying that non-myopic planning neither hinders nor helps in this case.
A.4. Environment Similarity
We evaluate ZILOT and the myopic baselines Pi+CIs and MPC+CIs on the walker environment (Tassa et al., 2018) in Table 3. Because this environment is very similar to the half aheetah environment we use in our main evaluation, we can reuse the same goal-space, tasks, and data collection method. These similarities are also visible in the performance of the three methods.
15

---

Zero-Shot Offline Imitation Learning via Optimal Transport
MPC+ ZILOT (ours)
0.72
0.20
0.64
GoalFraction
0.16
Wmin
0.56
0.12
0.48
MPC++Cls
ZILOT (ours)
0.08
0.40
16
32
64
8
16
32
64
Horizon
Horizon
Figure 8. Mean performance across five seeds in fetch slide large 2D for different planning horizons.
fetch_ Setch_pick.._. wwtp__icccccppppppeeeee and_place
fetch_push
Setch__lide__i_e www__dldddd____________
balfchestab
pointmaze, _miiium
1.00
1.00
1.00
1.00
1.00
0.75
0.75
0.75
0.75
0.50
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
0.00
ERrn RER
Pi+Cls MPC+CIs
ZILOT+h ZILOT+CIs
ZILOT+Unbalanced ZILOT (ours)
Figure 9. Single Goal Success Rate in the standard single goal tasks of the environments. We report the mean performance across 20 trials and standard deviationacross seeds.
Table 3. Evaluation the walker environment with the ha salfchhetab results repeated for comparison.
Task
W Wmin
GoalFraction f
Pi+CIs
MPC+CIs
ZILOT (ours)
Pi+Cls
MPC+CIs
ZILOT (ours)
walker-backflip
2.804110.056
1.73710.148
1.273:10.205
0.3410.07
0.89:1003
0.92:10.06
balker-backflip
3.039:10.292
2.444::0.189
1.709:10.093
0.49:00.07
0.70=10.08
0.81:10.99
walker-frontflipp
2.688:10.400
1.830:10.188
1.551
0.57:0.11
0.94:10.04
0.95:10.07
falkee--roont.lii-rnnn
2.5974:0.255
1.937--0.172
1.921=10.199
0.55+00.03
0.63
0.76:10.11
walker-hhop-backwardd
1.447+0.076
0.872:00032
0.836:10.100
0.64
0.78
0.84
wolker--hop-forward
0.932+10.098
0.663:10.071
0.467:1.044
0.95:10.05
0.9910.01
1.00:10.01
walker-ron-backwar
1.290410.144
1.050:100086
0.957-10.111
0.81-10.08
0.83+10.14
0.84+10.99
walker-sun-forwarrd
1.180
0.954
0.672:10.058
0.86+10.07
0.9710.03
0.99:1.001
walker-all
1.99710.047
1.436:1.0049
1.174:0.061
0.65:10.33
0.8410.02
0.89110.44
halffcheetsh--aackliip
3.089
4.28110.371
2.625:10.880
0.2811011
0.12:1.11
0.57
hal felf
2.879+00.427
3.044+00752
2.444
0.44+0.11
0.46++0.18
0.58++0.11
hallfcheet:ah=ffeefllip
1.544:10.127
1.695.10.147 147
1.295:10.094
0.7710009
0.79110.12
1.00
hal fihhhettt--rrrottiipp-uuuuun
1.3
2.083+40.104
1.955-+0.057
0.70++0.08
0.81++0.07
0.85++0.33
haa fchsettah--oo--baccee.
0,806++0.1
0.950+00.075
0.889++01107
0.96++0.03
0,90++0.02
0.96++0.03
hallfcheettah-hop-fooaaaa
1.58011.069
1.392:10.20
1.101:0.152
0.51110.07
0.62:10.14
0.58-10.12
hal Fochettah--un--bcckeerr
0.89740.092
0.679+-0.035
0.489-:0.167
0.96:+0.04
1.00++0.00
0.99+00.01
malffcheetth--ru-ffowwaa
0.857=10.04
0.8224002066
0.376:10.019
1.00:10.01
0.9410.08
1.00:10.00
halfcheetah-all
1.71740.101
1.868100.799
1.325:10.123
0.70
0.71
0.82-10.02
B. Forward-Backward Representations and Imitation Learning
In foundational paper in zero-shot, model-free II., (Pirotta et al., 2024) propose several methods based on the forward- backward (FB) framework (Touati & Ollivier, 2021). FB trains two functions F and B, which recover a low-rank approximation of the successor measure, as well as a parameterized policy (mz) These functions can be trained offline, without supervision, so that for each reward r, an optimal policy т., can recovered. This property gives rise to range of reward-bar based and accupancy-matching based methods for zero-shot II. In the following we will go over each method, and
16

---

Zero-Shot Offline Imitation Learning via Optimal Transport
discuss how differs from ZILOT in terms of objective. We will highlight how several methods do not directly apply setting, which involves expert demonstrations that are actionless, rough, and partial. We refer the reader to section C. for implementation details of baselines based on FB.
B.1. FB Imitation Learning Approaches
Behavioral Cloning The first approach in Pirotta et al. (2024) is is based gradient descent the latent find the policy T. that maximizes the likelihood of expert actions. Since this approach strictly requires expert actions it does not apply in our case.
Neward-Based Imitation Learning (Pirotta et al., 2024) derive two reward-based zero-shot IL methods maximizing the reward r(.) = pp()))/) (p)ppa(() (ERra) (Ma al., 2022; Kim et al., 2022) and its regularized counterpart r()) p ()/(oP())) ())pp(() (-)) (RERFa) (Reddy et al., 2020; Zolna et al., 2020). While ZILOT's objective is based on a Wasserstein distance, these rewards are derived from regularized f-divergence objectives. These objectives are fortunately tractable, and can be minimized by solving an RL problem with additive rewards. In practice, this corresponds to assigning a scalar reward to each state visited by the expert, without considering the order of the states in the expert trajectory. However, as stated in Section 4.2 of Pirotta et al. (2024), this regularization comes at cost, particularly if the state does not contain dynamical information, or ergodic MDPs. In this case, a policy can maximize the reward by remaining in the most likely expert state, and the objective might be optimized by degenerate solution. On the other hand, such solution would be discarded by ZILOT, which uses unregularized objective.
Nonetheless, these two instantiations are fully compatible with partial and rough demonstrations. Thus, we provide an empirical comparison in Section 5.
Distribution Matching A further approach in Pirotta et al. (2024) finds the policy л. whose occupancy matches the expert occupancy w.r.t. different distances the space of measures. ZILOT also performs occupancy matching, but with respect to Wasserstein distances. However, ZILOT is designed to handle state abstraction, i.e. partial states. To the best of our understanding, distribution- and feature-matching flavors of FB-IL require the demonstration to contain full states, unless further FB representations are trained to approximate successor measures over abstract states. While the standard implementation of iistribution-macchnng FB-IL cannot imitate rough demonstrations, we believe that an extension in this direction may be interesting for future work.
Goal-Based Imitation Pirotta et al. (2024) also instantiate a hierarchical, goal-based imitation method, in which the FB framework is only used for goal-reaching. This idea is closely related with one of our baselines (Pi+CIs). However, their framework assumes that trajectories to imitate are not rough and, instead of using a classifier, the goal can be chosen at a fixed offset of in time for each time-step. In any case, their approach remains myopic as per Proposition 3.1. Empirically, Pirotta et al. (2024) observe that this instantiation of FB-IL does not significantly outperform an equivalent method relying on TD3+HER instead. As the latter method is very similar to our Pi+Cls baseline, do not investigate this approach further in this work.
C. Implementation Details
C.1. ZILOT
The proposed method motivated and explained in section 4. We now present additional details.
Sinkhorn
First, we rescale the matrix C' by Tmax and clamp it to the range [0, 1] before running Sinkhorms algorithm. The precise operation performed is
C min (1, ma C/Tmax))
(19)
This is done so that the same entropy regularization e 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., 2021) and (Cuturi et al., 2022). We run our Sinkhorn algorithm for r 500 iterations with a regularization factor of 0.02.
17

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Truncation When the agent gets close to the end of the expert trajectory, then we might have that lx k + H. i.e. the horizon is larger than needed. We thus truncate the planning horizon the estimated remaining number of steps (and least 1), L.e. we set
Hactual max (1, min(tx - k. H))
(20)
Unbalanced OT As mentioned in the main text in section 5.3, we can use unbalanced OT (Liero et al., 2018; Séjourné et al., 2019) to address that fact that the uniform marginal for the goal occupancy approximation may not be feasible. Unbalanced OT replaces this hard constraint of T 1N = 1M into the term (.KL(TT -l,, 1M) in the objective function. For our experiments we have chosen Çb- 1.
C.2. TD-MPC2 Modifications
As TD-MPC2 (Hansen et al., 2024) is already a multi-task algorithm that is conditioned on learned task embedding & from task id i, we only have to switch out this conditioning to a goal latent 2g to arrive at a goal-conditioned algorithm as detailed in table 4. 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 .M (assuming truncation after goal reaching, see section 2.3). For training we adopt all TD-MPC2 hyperparameters directly (see table 7). As mentioned in the main text, we also train small MLP to predict W that regresses on V.
Table 4. Our modifications to TD-MPC2 to making it goal- instead of lask-conditioned.
TD-MPC2 (Hansen et al., 2024) "GC"-TD-MPC' (our changes)
Task/Goal Embedding
t= E()
2g hg(g)
Encoder
h(s
2 hs))
Dynamics
z f(z, a,t)
z' = f a)
Reward Prediction
r = R(z, 0,1)
r = R(z, O, Zg)
Q-function
= a(2,a,t)
= Q(2, a, Zg)
Policy
a (22
T 2g)
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 = ( g) = (Q(z, π T(, =g); 2g) where x = h(ss,, 2g = (g). To speed-up computation, we train separate network that estimates the value function directly. employs a two-stream architecture (Schaul et al., 2015; Eysenbach et al., 2022) of the form V (7 (g) = o(z) (22g) where Φ and eb are small MLPs for fast inference of pair-wise distances.
(s,g)
Our GC-TD-MPC2 is trained like the original TD-MPC2 with two losses additionally employing HER (Andrychowicz et al., 2017) to sample goals which we discuss in detail in Appendix C.4. The first loss combines a multi-step loss for d and h with a single-step TD-step loss for R and Q
H
L = E (...,")...HD D So Pa.nHER (1)
A - sg(h()))) + CB(R(2), ar zg. )..) + CE(Q(2), at, Z zg.). 9:)) Σ A=
(21)
where sg is the "stop-gradient"-oprreee -operator, Zz, Age. and 2/ are defined in Table 4, rewards are r. = Il.. 11, and (undiscounted) TD-targets are q: max(r, + llaeton Q(1,, π (=bs Zg.) )g.), -Tmax). The second loss is a SAC-style loss for ग (Haarnoja et al., 2018)
H Σ
Cx = E (s.a,""oonn0) D
(aQ(z), zm.) (- 2gx3) 2t = d(ze, az), 20 = h(50).
(22)
90.8 HER (50.8)
Additional to GC-TD-MPC2 we then also train our goal-conditioned value function V* =g) =b=)= (x(g)))))))))) using the same TD-targets as in Equation (21)
H X (é(zz) (x(z))))))))))))))))))())))))( =t)²
L
E Sp
Σ t
(23)
90.n THERR (50.))
18

---

Zero-Shot Offline Imitation Learning via Optimal Transport
C.3. Runtime
ZILOT runs at 4Hz on an Nvidia RTX 4090 GPU, depending on the size of H and the size of the OT problem. Given that the MPC+CIs method runs at around 25 to 72Hz with the same networks and on the same hardware, it is clear that most computation is spent on preparing the cost-matrix C' and running the Sinkhorn solver. Several further steps could be taken to speed-up the Sinkhorn algorithm itself, including 77-schedules and/or Anderson acceleration (Cuturi et al., 2022) as well as warn-starting it with potentials, e.g. from previous (optimizer) steps or from trained network (Amos et al., 2023).
C.4. Goal Sampling
As mentioned in the main text, we follow prior work (Andrychowicz et al., 2017; Bagatella & Martius, 2023; Tian et al., 2021) and sample goals from the future part of trajectories in D3 in order to synthesize rewards without supervision. The exact procedure is as follows:
With probability Piuture 0.6 we sample a goal from the future part of the trajectory with time offset lo Geom (1-)..
With probability Prext 0.2 we sample the next goal in the trajectory.
With probability Prand 0.2 sample a random goal from the dataset.
C.5. Training
We train our version of TD-MPC2 offline with the datasets detailed in table 5 for 600k steps. Training took about 8 to hours on single Nvidia A100 GPU. Note that TD-MPC2 samples batches of transitions per element, we effectively sample 3 256 768 transitions per batch. The resulting models are then used for all planners and experiments.
Table 5. Environment description. We detail the datasets used for training.
Environment
Dataset
# Transitions
fetch.push
WGCSL (Yang et al., 2022) (expert+random)
400k + 400k
WGCSL (Yang et al., 2022) (expert+random)
400k + 400k
Beechhslide.larre 2D halfcheetah
custom (curious exploration (Pathak al., 2019))
500k
custom (curious exploration (Pathak et al., 2019))
500k
pointmaze_medimn
D4RL (Fu et al., 2021) (expert)
1M
C.6. Environments
We provide environment details in table 6. Note that while we consider an undiscounted setting, we specify y for the goal sampling procedure above.
Table 6. Environment details. We detail the goal abstraction é, metric h, threshold e, discount factor y used for each environment.
horizoo H, maximum episode length Tmax, and
Environment
Goal Abstraction
Metric h
Threshold c
Horizon H
Tmax
Y
fetch.push
(2, y, eeube
. |/2
0.05
16
50
0.975
fetich pi and pl plece fetch.slide.large 2D halfcheetah pointmaze.medium
(x, y, z) cuub
: ||2
0.05
16
50
0.975
(x, y, z cuube
112
0.05
16
50
0.975
(x, By)
||. l|2 0.50
32
200
0.990
(x, y)
l|2
0.45
64
600
0.995
The environments fetch.push and halfcheetah is not goal-conditioned by default, we define our own goal range be (z, 0,)) E [-5, 5] x [-47, 4 fetch_s 1 idde _argeel is variation of the fetch_s 1 ide environment where the table size exceeds the arm's range and the arm is restricted to two-dimensional movement touching the table.
and www.ntmaaeemmdddnn are used as is. As
Note that the ha 1 fcheetah environment does not reduce with any kind of modular operation, i.e. states with 0 and 2m are distinct.
19

---

Zero-Shot Offline Imitation Learning via Optimal Transport
C.7. Tasks
The tasks for the fetch and pointmaze environments are specified in the environments normal goal-space. Their shapes be seen in the figures in appendix E. To make the tasks for halfcheetah more clear, visualize executions of method in the figures 10, 11, 12, 13, 14, and 15.
1n
O1
-2T
-3m7
0.0
0.5 Trajectory
1.0
1.5
2.0
2.5
3.0
3.5
4.0
x
Goal
X
Figure 10. Example trajectory of ZILOT (ours) in hal taaTfchheetchhhaaff1-------m
1s 0T 1 TT
3m
-0.2 Trajectory
-0.1
0.0
0.1
0.2
X
Goal
x
Figure 11. Example trajectory of ZILOT (ours) in halfcheetah-bacckil.
27
-1m
2
4
5
6
x
Goal
Trajectory
x
Figure 12. Example trajectory of ZILOT (ours) in hal
20

---

Zero-Shot Offline Imitation Learning via Optimal Transport
1n
0m
X
x
1m
-2
T
2
4
x
Goal
Trajectory
X
Figure 15. Example trajectory of ZILOT (ours) in hallfcheetah-ho--ooooaa
311
217
X
0T
-1m
0.0
0.5 Trajectory
1.0
1.5
2.0
2.5
x Goal
x
Figure 13, Example trajectory of ZILOT (ours) in fahh
1n
x
0m
-1n
-5
-4
-3
-2
-1
X
Goal
Trajectory
X
Figure 14, Example trajectory of ZILOT (ours) in ballfchhetth---o--accaaaa
C.8. Task Difficulty
This section investigates the ability of ZILOT to imitate trajectories that do not appear in the offline dataset it is trained on. As ZILOT uses learned dynamics model and an off-policy value function, it should in theory be able to stitch together any number of trajectories in the dataset. To get some qualitative intuition we overlay the following: first, a kernel density estimate of the data distribution the offline datasets, second, an expert trajectory to imitate, and finally the five trajectories that are closest to the expert w.r.t. the Wasserstein distance under the goal-metric We present a few tasks for each environment in Figures 16, 18, 19, 17, and 20.
21

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Comparing the density estimates and the expert trajectories, we bution. Although, especially in ha I fcheetah, there some tasks, such hop-forwas rd and back faccklipp-2e -il-rrunnnt ng with very little coverage which might explain the bad performance of all planners in these tasks (see table 1). Comparing the selected trajectories with the expert trajectory, it is also evident that the expert demonstrations are not directly present in the datasets. Thus, ZILOT' is capable of imitating unseen sequences of states, as long as each individual state within the support of the training data. other words, ZILOT capable off-policy learning, or trajectory stitching.
that essentially all expert trajectories are within distri-
(a) hww-hoo-oooree forwadd
(b) hap-backwwrd
(c) backf lick
(d) frontflip
(e) frontilip-running
( fallrrrrpn
Figure 16. The 5 trajectories (blue) from the dataset that are closest to the expert trajectory in different ha halffan fcheethh tasks (orange) overlayed over a kernel density estimate of goal occupancy in full training dataset.
22

---

Zero-Shot Offline Imitation Learning via Optimal Transport
(a) L-dense
(b) U-dense
(c) S-dense
Figure 17. The 5 trajectories (blue) from the dataset that closest the expert trajectory in different wwwch.slidc.aarge...o.. tasks
(orange) overlayed over kernel density estimate of the goal occupancy in the full training dataset.
(a) l.-dense
(b) U-dense
(c) dense
Figure 18. The 5 trajectories (blue) from the dataset that closest the expert trajectory in different fetch.push tasks (orange) overlayed over a kemee density estimate of the goal occupancy in the full training dataset.
(a) L.-dense
(b) U-dense
(c) S-dense
Figure 19. The 5 trajectories (blue) from the dataset that are closest to the expert trajectory in different fetch pick and place tasks (orange) overlayed over kernel density estimate the goal occupancy the full training dataset.
23

---

Zero-Shot Offline Imitation Learning via Optimal Transport
(a) circeedeense
(b) path-dense
Figure 20. The 5 trajectories (blue) from the dataset that closest the expert trajectory in different pointmaze mediim tasks (orange) overlayed over kernel density estimate of the goal occupancy in the full training datasct.
C.9. Hyperparameters
Table 7. TD-MPC2 Hyperparameters. We have adopted these unchanged from Hansen al. (2024)
Name
Value
Name
Value
1r
3e-4
num_bins
101
batch-size
256
vmin
-10
n.steps ("horizon")
3
vmax
10
rho
0.5
num.enc.layers
grad.clip.norm
20
enc.dim
256
enc.lr_scale
0.3
nun.channels
32
value.coef
0.1
mlp_dim
512
reward.coef
0.1
latent_dim
512
consistencyycooff
20
bin.dim
12
tau
0.01
num_q
5
log_std_min
-10
dropout
0.01
1og.std.max
2
simnorm_dim
entropy coef
le-4
24

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Table 8. Hyperparameters used for iCEM (Pinneri al., 2020). We the implementation from Pineda (2021).
(a) ICEM hyperparameters for all MPC planners.
(b) ICEM hyperparameters for curious exploration.
Name
Value
Name
Value
num iterations
num iterations
population.size
512
population.size
512
elite.ratio
0.01
elite.ratio population decay factor
0.02
population decay factor
1.0
0.5
colereednoise.xxonnett
2.0
coloredinoise..xponenn
2.0
keep.elite_frac al alpha
1.0
keep.elite.frac
1.0
0.1
a alpha
0.1
horizon
20
C.10. FB Implementation Details
Since there is implementation available for FB-IL directly, have adopted the code for FB (Touati & Ollivier, 2021) according to the architectural details in appendix D.3 and the hyperparameters appendix D.4 of FB-II. (Pirotta al., 2024). The main architectural changes consisted of changing the state input of the B networks to only goal input, suggested in (Touati & Ollivier, 2021) as well as adding last layer in the B networks for L2 projection, batch normalization, or nothing, depending on the environment.
follow the specifications of Pirotta et al. (2024) whenever possible. As hal Salfcheetah and maze are also used in their evaluations we have adopted their hyperparameters for these environments. For our fetch environments, we used the hyperparameters most common in the environments except for the discount 7 which we adjusted 0.95 to account for the shorter episode length. Finally, we have found that the FB framework seems to be ill-adjusted be trained magnitude less data then in the original experiments (Touati & Ollivier, 2021; Pirotta et al., 2024). For some environments, performance started to deteriorate rather quickly, so we to report the best performance encountered during training when evaluating every 50k steps (see D.2). We provide the full set of hyperparameters in table 9.
order of
Table 9, Hyperparameters used for FB-IL 1B-IL training. Closely follows table I in appendix D.4 of (Pirotta et al., 2024) for hal fcheetah naze.
Environment
fetch
halfcheetah
maze
Representation dimension
50
50
100
Batch size
2048
2048
1024
Discount factor 7
0.95
0.98
0.99
Optimizer
Adam
Adam
Adam
learning rate of learning rate of B learning rate of π Normalization of B Momentum for target networks Stddev for policy smoothing Truncation level for policy smoothing
10-4
10-4
10-4 4
10-4
10
10- 6
10-4
10-1
10-6
L2
None
Batchnorm
0.99
0.99
0.99
0.2
0.2
0.2
0.3
0.3
0.3
Regularization weight for orthonormality
Numer of training steps
2 108
2- 10
2 108
25

---

Zero-Shot Offline Imitation Learning via Optimal Transport
D. Hyperparameter Searches
D.I. Classifier Threshold
As mentioned in the main text, we perform an extensive hyperparameter search for the threshold value of the goal classifier (CIs) for the myopic methods Pi+CIs and MPC+CIs as well as for the ablation of our method ZILOT+CIs. In figures 22 and 21 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 all tasks attain maximum performance with the same threshold value showing that this hyperparameter is rather hard to tune.
0.20
GoafFraction
Goaliraztion
Wain
0.15
Wwin
0.4
0.10
0.4
1
2
2
3
4
Threshold
Threshold
Threshold
Threshold
S-sparse
U-dense
S-sparse
U-dense
Best Goaliraction
Lddense
U-sparse
Best Goaliraccion
L-dense
U.sprr
S-dense
L-sparse
S-dende
L-sparse
(a) fet weccch.pic..nd...p....
(b) fet ch-pus
0.8
GoafFraction
GoalFrractio
Wean
0.15
2-
0.10
0.2
1
2
3
Threshold
Threshold
S-sparse
U-dense
Threshold
Threshold
Best GoalFraction
Ldenne
U-sparse
Best Wais
path dense
circle dense dersle-dense
S-dense
L-sparse
Best GoalFractian
path-sparse
circle-sparse
(c) fetch_s
iide_d_111 arge_2D
(d) pointmaze.med um
Goaliraction
0.6
3
Threshold
Threshold
hop forward
backflip running
Best CoalFraction
hop-backward
run-forward
frontflip
backflip
run-back uun-backward
frontflip running
(e) halfcheetah
Figure 21. ZCL hyperparameter search for Cls threshold.
26

---

Zero-Shot Offline Imitation Learning via Optimal Transport
0.30
0.25
0.8
0.25
0.20
GoalFraction
0.6
Wrin
0.6
Weon
0.15
0.4
0.15
0.10
0.2
0.10
0.2
1
2
5
2
3
4
Threshold
Threshold
Threshold
Threshold
Were
S-sparse
U-dense
Best Wri
S-sparse
U-dense
Best GoalFraction
L-dense
U-sparse
Best Goaliraction
L-dense
U-sparse
S-dense
L-sparse
S-dense
L-sparse
(a) Pi+Cls for fetch_pi ck_and_pll plcce
(b) MPC+CIs for fetch_pi ck_and_placc
0.25
0.8
0.8
Wei
0.20
0.15
GoalFactoon
0.6
0.4
Wer
0.20
0.15
0.6
0.10
0.2
0.05
1
2
5
3
1
2
3
Threshold
Threshold
Threshold
Threshold
Best
sparre
U.dense
Best Wenn
S-sprree
U.dense
Best GoalFraction
L-dense
U-sparse
Best Goaliraction
L-dende
U-sparse
S-dense
L-sparse
S-dense
L-sparse
(c) Pi+Cls for fetch. push
(d) fetch fetch.push
0.35
0.30
0.30
Goalfraction
0.4
0.25
Goaliraction
0.6
0.25
0.20
0.2
0.20
0.15
0.15
2
Threshold
Threshold
Threshold
Threshold
Best Wrn
S-sparse L-dense
U-dense
Best Wetn
S-sparse
U-dense
Best GoalFraction
U-sparse
Best Goaliraction
L-dense
U.sparse
S-dense
L-sparse
S-dense
L-sparse
(c) Pi+CIs for fetch_s ddd_ rge_21
(f) MPC+CIs for fetch_s 1 de_ 1 arge_211>
1.0
1.00
GoalFraction
GoalFraction
0.75
0.6
0.4
0.25
0.2
Threshold
Threshold
Threshold
Threshold
Best Wain
hop forward
backflip-runme
Best Wein
hop hor-forward
backflip running
Best Goalfraction
hop-badkward
run-forward
Best Gealfraction
hop-backward
run-forward
trontflip
backflip
run-backward
frontflip
backflip
run-backward
frontilip-ruming
frontfllp-running
(g) Pi+Cls for ha 1 f fchestah
(h) MPC+Cls for ha 1 acheetah
1.0
1.0
Coairraction
0.8
Foatoon
0.8
2
0.6
0.4
0.4
0.2
0.2
1
2
1
2
Threshold
Threshold
Threshold
Threshold
Best Wan
path-dense
drcie-dense
Best Won
path-dense path-sparse
drcie-dense
Best GoalFraction
path-sparse
dicce-sparse
Best GoalFraction
dicce-sparse
(i) Pi+Cls for mointmaze.mediiin
( MPC+CIs for pointmaze.medium
Figure 22, Pi+Cls and MPC+CIs hyperparameter searches for Cls threshold in each environment.
27

---

Zero-Shot Offline Imitation Learning via Optimal Transport
D.2. FB-IL Training steps.
As mentioned in section C.10 we report the best evaluation results of all FB methods that occur during training. In figures 23
and 24 report the evaluation performance for different training lengths.
0.250
0.275
2.0
0.4
0.225
GoalFraction
0.250
GoalFraction
m
Wmin
Wmin
0.3
0.200
0.225
1.5
0.175
0.200
0.2
1.0
0.150
1
2
1
2
0
1
2
0
1
2
Step
1e6
Step
1e6
Step
1e6
Step
1e6
(a) fetch pick and place
(b) pot nnzmaee medium
0.32
0.26
0.6
0.30
0.24
2.5
GoalFraction
GoalFraction
Wmin
0.28
0.22
Wmin
0.4
2.0
0.26
0.20
0.18
1.5
0.2
0.24
2
Step
1e6
Step
1e6
Step
1e6
Step
1e6
(c) fetch slice large 2D
(d) halfcheetah
0.30
0.250
m
GoalFraction
Wmin
0.225
0.25
0.200
m
0.20
0.175
Step
1e6
Step
1e6
(c) fetch.push
Figure 23, ER evaluation performance during training.
28

---

Zero-Shot Offline Imitation Learning via Optimal Transport
0.250
0.275
2.0
0.4
GoalFraction
GoalFraction
0.225
0.250
Wmin
Wmin
0.3
0.200
0.225
1.5
0.175
0.200
0.2
1.0
0.150
0.175
0
1
0
1
2
1
0
1
2
Step
1e6
Step
1e6
Step
1e6
Step
1e6
(a) fetch..iick.and.plac
(b) www.n//maammmdiiiinn
0.32
0.26
0.6
0.30
0.24
2.5
GoalFraction
MW
GoalFraction
Wmin
0.22
Wmin
0.4
2.0
0.26
0.20
0.18
1.5
0.2
0.24
2
Step
1e6
Step
1e6
Step
1e6
Step
1e6
(c) www.hhhsiieeeeooeoeeef war.ee.eee.lrr.....
(d) halfcheetah
0.30
0.250
GoaiFraction
Wmin
0.225
0.25
0.200
a
0.20
0.175
Step
1e6
Step
1e6
(e) tetch.push
Figure 24, RER evaluation performance during training.
29

---

Zero-Shot Offline Imitation Learning via Optimal Transport
E. 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 ch environments display some natural symmetries, decided split evaluations between all four symmetrical versions of them. Further, quickly want stress that these trajectories are shown in goal-space. This means that if the cube in fetch is not touched, as is the 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.
Hirds
ZLOTT loura Unzslanced
28OT
ZIOT-A
ZROTICk
neDs
ZRoT+Cs
Z017 toural Unpelaacee ZILOT-A
ZIOT lano ZILOT-Unsslannn
ZLOT-A
ZBOT4Ca
NIC+
ZIOT (ound) ZILOT Uncelanced
ZBOT-A
ZBOT
RieOs
NFCeCis
ZILTT (ours) 2ioohnen
PleCtx
NFCects
ZILT FloTtwhnattnned
Z1LTTHh
ZROT+Ce
(a) U-dense
(b) U-sparse
Figure 25.
30

---

Zero-Shot Offline Imitation Learning via Optimal Transport
toural Unidaanced ZiOT
2LD foural zuorr Unisdaned ZIOT-A
erсcca
ZLOT (ers) ZUoT+Umpatinet
ZiaTen
ZBOT+Ca
2+-
NrCaca
ZLDT (ours ZlorrUppallneed
ZILOTA
ZEOT<ch
NrCaci
ZI1T (our)) ZEoTsUnaakeee
NiCcch
ZIOT PioTTthrcaeeeed
ZELTTA
7R0T4ck
(a) L-sparse
(b) L-dense
Figure 26. fetch.picckand.placs
ZLOT tonddnenedeeee
N:0s
NICics
2107 (pur)) Unnauncc
ZLOTA
NiCs
NICicb
ZLOT toural ZIOT Unsulaned
ZIOT-A
2226
N42x
ZIST (ond) Undtenas ZEOTed
ZL0T4Ch
NrCkCk
ZLDT 2LLTT&hncaeeeeed
ZEOTTA
ZROT+Ch
sSSSS
(a) S S-dense
(b) S-sparse
Figure 27. fetch pick and place
31

---

Pero-Shot Offline Imitation Learning via Optimal Transport
R+DD
NiCeCls
ZILT (ourd) ZLOT-Unaatnons
ZLOted
NFC+Cle
ZILTT tourt Tiootelncatonnnd
ZILOTed
ZLOT+Ch
H+0s
NFC+Cls
ZIOTT (ours) ZLoT+Uhabboedd 2EST+A
20.0+C0
NFCecls
DLOT tourd 21oretantteeed
ZILOTH
2ROT+CB
a+
foursi UusThthisseeed
28OT+ck
ZuoT+Unsatanee
ZIOT-h
ZuoTice
nets
(oural Unnatlncedd
ZBOT
(a) U-dense
(b) U- sparse
Figure 28. fetch largee
ZLOT loura) Unielanced
20DT ZUOT Unselanced
ZILTT
ZLOTICL
744h
NrCaCh
ZIOT berd ZLOT Unbeleenndd
ZR0T+A
ZLOT+CB
NTC4
ZLOT (ours) 2loorUnneeeeeee
ZBOT-A
ZBOT4ch
NTCaCia
ZILT (ourd) ZILOT Unzslanced ZEOTTA
Z8OT+Cl
NFci
ZILTT 20oor-Unntteneee ZLLTTh
Z4OT4Ch
(a) L-sparse
(b) L-dense
Figure 29. fetch_alide 1e__gieeee_rarr@
32

---

Zero-Shot Offline Imitation Learning via Optimal Transport
Zioretnastbned
N:D
erCICs
ZloT ZLOT Unedeneed Z
ILOTICS
hiCs
NICich
ZIDT toura) ZILTT Unradaaced ZROT-A
ZEOTICb
7+
NrCaCa
(sSn)) Unotenne ZLOT
ZLOT+Ch
C+
NtC+
ZLDT ZLoT+Uhhatenced
ZILOT-h
ZBOT+ch
(a) S-dense
(b) S-sparse
Figure 30. fetch
large 21
ZLOT loura) Unialanced ZLOTIA
2LOT ZLOTT Unselanced
ZILOT-A
nets
ercach
ZLOT berd Unbreeneeeee
ZLOT+C5
Nrc4ch
ZIOT toursl ZIL0T Uncollaced
ZBLTTA
ZBOT+ch
NFCacia
ZI57 (ound) Unstenedd ZLOT-A
28.T+Cl
NFC+
ZILOT 2iooo-Uhaatrneed ZILOT-
Z4OT4Ck
(a) U-denso
(b) U- U-sparse
Figure 31. fetch_ ppuhh
33

---

Zero-Shot Offline Imitation Learning via Optimal Transport
7l+Ds
NFC+Ch
ZIST bourd ZLoT+Unpaltnees
ZLOT+A
ZLOT+CN
THC
NFC+Ch
ZLDT (ours) ZlooteUnaslence
ZIOTT
ZLOT+Ch
Ri-Gs
NFCacis
ZIOT lourt) ZioTTsnnahheedd Z1.T+a
Z8.OTChh
NFCacK
ZILTT ZloTsthncateneed
ZiLTT-h
Z4oT<ch
toursi ZU0T+Umabboed
NPCICS
ZIOT loural ZROT Whiseeeced
ZEOTICS
HEDS
NrCIch
2101 loura 210 HOrodaace
ZEOT-
ZLOTICE
(a) L-sparse
(b) L-dense
Figure 32, fetch.push
(ours)
Heds
NiCiCs
ZLOT lours ZEOT Uncabnee ZKG
Hebss
Nrcrcc
ZLDI loursi Z1OT #Uroslanced
ZROTTA
ZuoF
NrCiCs
ZIST (sua) ZLOT Unaalnndd ZLOT-A
NiCs
NrCICt
ZLDT tours) Z1OT Uncalanced
ZIOT-A
ZeoTick
s3s6ss
S
(a) S dense
(b) S- S-sparse
Figure 33. fetch.push
34

---

Zero-Shot Offline Imitation Learning via Optimal Transport
R+ch
MPC+CS
ZILOT Tours)
ZILOT+A
ZILOT
(a) ci rccc circeesparss
tours)
(b) path-sparse
8+Cb
MPC+CS
ZIDT fours)
ZLOTAUrbalancer
ZILOT+A
ZIOT+CS
(c) path-dense
RiCh
Cls
ZILOT tours)
ZILOT Unbalancos
ZILOT
ZILOT:C5
(d) circle-dense
Figure 34. pointmaze medium
35

---

Zero-Shot Offline Imitation Learning via Optimal Transport
CiC
MPC Cls
ZILOT tours)
ZILOT Unbalanced
ZILOT
ZILOT
(a) run- forwaroaaa rd
R+Cbs
MPC+Cs
ZILOT lours)
ZILOT #LOTluubblancod
ZILOT+6
ZILOT+C5
(b) run run-backward
C
ZILOT
ZILOT Unballanced
2
(c) backflip
R+C
MPC+Cls
ZILOT lours)
aLot+Unbalanced
ZILOT+A
ZIUO+CC
(d)
Figure 35. hal fcheetah part 1
36

---

Zero-Shot Offline Imitation Learning via Optimal Transport
R+Os
MPC+CB
ZILOT tours)
ZILOT Unbblaanced
ZILOT+A
ZIUO++G
(a) hop-backward
A+C
MPC+Css
ZILOT
ZILOT #UoT+&nbbttacess
ZIUO++b
ZILOT
(b) front frontflip-running
ZILOT fours)
Uinallaaaa
ZILOT4A
(c) backflip-runniiee
8+Ch
MPC+CS
ZILOT fours)
ZILTT
ZILOT4h
ZIUO++C
(d) frontflip
Figure 36. hal fcheetah part
37


Reviews:
[{"review_id": "FMXWHb8a3O", "reviewer": "Reviewer_8Rj6", "summary": "The problem of greediness in goal conditioned imitation is resolved by matching goal occupancies. The proposed algorithm can learn from a single demonstration with partial observability. It is shown that minimizing wassertein distance between goal occupancies of expert and learner is equivalent to minimizing an optimal transport objective. This is utilized to learn policies. The proposed setup is compared to existing works with thorough experiments by measuring the proportion of goals achieved by each algorithm.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their feedback and comments.\n\nThank 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.\n\nThe 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. \nWe 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.\n\nWe remain available for further discussion for the rest of the rebuttal period.", "claims_and_evidence": "- ZILOT is non-myopic. Evidence theoretical from a pathological example where some other algorithms might fail, and empirically from certain environments where there are multiple goals to be achieved in a single trajectory.\n- ZILOT giving good Offline and Zero-Shot performance and the approximations being enough to learn a good policy demonstrated via experiments.\n- Ablation studies demonstrate the usefulness of some individual components.", "methods_and_evaluation_criteria": "Yes. Thorough experimentation. Comparison to relevant work.", "theoretical_claims": "Not in depth. No issues from surface level reading.", "experimental_designs_or_analyses": "Did not run any code. The design description in paper and appendix seems very sound.", "relation_to_broader_scientific_literature": "Zero Shot Learning that is not Greedy.", "essential_references_not_discussed": "none that I am aware of", "other_strengths_and_weaknesses": "Strengths:\n- Very detailed results and diagrams showing how proposed algorithm works in comparison to others. One result in main paper, rest in appendix.", "other_comments_or_suggestions": "- Acronym GC-RL is used without first introducing it.\n- I don't understand why there is a subscript 1 in definitions of $\\mathcal{D}_\\beta$ and $\\mathcal{D}_E$ . Line 78 col 2.", "questions_for_authors": "none", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GhsJA4SAnG", "reviewer": "Reviewer_LRC1", "summary": "- The paper concerns zero-shot offline model-based imitation learning\n\n - A new method is proposed based on Optimal Transport to match the occupancy measure of the learning and expert policies\n\n - Experiments conducted on three tasks (Fetch, HalfCheetah, PointMaze), comparing the proposed method with existing baselines", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your assessment and valuable feedback. We first address your main concern, our empirical evaluation, and then your other questions.\n\n# Q5 on further experiments\n\nOur 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.\nWe 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:\n\n|Task|$W_{\\min}$ $\\downarrow$|||GoalFraction $\\uparrow$|||\n|-|-|-|-|-|-|-|\n||Pi+Cls|MPC+Cls|ZILOT (ours)|Pi+Cls|MPC+Cls|ZILOT (ours)|\n||||||||\n|walker-backflip|2.804±0.056|1.737±0.146|**1.273±0.205**|0.34±0.07|**0.89±0.03**|**0.92±0.06**|\n|walker-backflip-running|3.039±0.292|2.444±0.189|**1.709±0.093**|0.49±0.07|0.70±0.08|**0.81±0.09**|\n|walker-frontflip|2.688±0.400|1.830±0.185|**1.551±0.086**|0.57±0.16|**0.94±0.04**|**0.95±0.07**|\n|walker-frontflip-running|2.597±0.265|**1.937±0.172**|**1.921±0.149**|0.55±0.03|**0.63±0.16**|**0.76±0.11**|\n|walker-hop-backward|1.447±0.076|**0.872±0.032**|**0.836±0.100**|0.64±0.11|**0.78±0.06**|**0.84±0.07**|\n|walker-hop-forward|0.932±0.098|0.663±0.071|**0.467±0.044**|0.95±0.05|**0.99±0.01**|**1.00±0.01**|\n|walker-run-backward|1.290±0.148|**1.050±0.086**|**0.957±0.111**|**0.81±0.08**|**0.83±0.14**|**0.84±0.09**|\n|walker-run-forward|1.180±0.105|0.954±0.079|**0.672±0.058**|0.86±0.07|**0.97±0.03**|**0.99±0.01**|\n||||||||\n|walker-all|1.997±0.047|1.436±0.049|**1.174±0.061**|0.65±0.03|0.84±0.02|**0.89±0.04**|\n||||||||\n\nOur 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.\n\n# Q1 model-free IL adaptation\n\nTo 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.\n\n# Q2 on OT vs. other distances\n\nThe 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).\n\n# Q3 on finite horizon planning\n\nWe 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).\n\n\n# Q4 on sampling errors\n\nEstimation 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. \n\n# Q6 on the comparison to IL and IRL\n\nansewredThe 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.\n\n# Q7 on scalability\n\nZILOT 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.\n\n---\nWe 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.\n\nReferences:\n1. Pirotta et. al. Fast Imitation via Behavior Foundation Models. ICLR ‘24\n2. Pathak et. al. Zero Shot Visual Imitation. ICLR ‘18", "claims_and_evidence": "- The main claim is that the proposed method optimizes an occupancy matching objective using Optimal Transport. In terms of methodological contribution, the approach appears novel. Using Optimal Transport for occupancy matching is a sound idea.\n\n- The paper claims to address a zero-shot imitation learning problem, but it is restricted to model-based, finite-horizon MDPs. The broader literature on imitation learning includes several efficient model-free algorithms, which limits the contribution of this work.\n\n- The experiments are limited and insufficient to fully support the main claims. Only three tasks are considered, while MuJoCo itself contains multiple tasks, yet the authors only evaluate on HalfCheetah.", "methods_and_evaluation_criteria": "- The method of using Optimal Transport for occupancy matching is reasonable.\n\n- Besides Optimal Transport, are there alternative ways to optimize occupancy matching, such as using KL divergence? Can the authors comment on this?\n\n- The main problem formulation is restricted to finite-horizon MDPs. Can the proposed method be extended to infinite-horizon models?\n\n- The occupancy estimation relies on sampling, which can introduce high errors and impracticalities. Additionally, the method requires dynamics estimation, which demands significant data and introduces further estimation errors.\n\n- The proposed method is specifically tailored to model-based settings. Can it be extended to model-free settings, which are often more practical?\n\n- As mentioned, the experimental evaluation is **weak**—only three simple tasks are considered. Why were other MuJoCo tasks not utilized for evaluation?", "theoretical_claims": "The paper provides minimal theoretical claims and proofs.", "experimental_designs_or_analyses": "The experimental setup is reasonable, but as previously mentioned, the experiments themselves are weak. The authors should consider extending the study with more challenging tasks typically used in prior imitation learning research.", "relation_to_broader_scientific_literature": "The contributions relate to reinforcement learning, imitation learning, zero-shot learning, and optimal transport. The approach has potential applications in fields where imitation learning is valuable, such as healthcare and autonomous driving, particularly when datasets are limited", "essential_references_not_discussed": "The paper should discuss more relevant work on model-free imitation learning and inverse reinforcement learning. These could suggest alternative approaches for handling model-free, infinite-horizon models and different ways to optimize occupancy matching.", "other_strengths_and_weaknesses": "- The approach is sound, and the use of Optimal Transport is an interesting direction.\n\n - The primary weakness is the experimental evaluation. This undermines the credibility of the claims and makes the paper not yet ready for publication.", "other_comments_or_suggestions": "Please see my questions", "questions_for_authors": "- Can the authors clarify why the proposed method is restricted to model-based settings? Could it be adapted for model-free imitation learning?\n\n- What are the main advantages of using Optimal Transport over other occupancy matching techniques, such as KL divergence?\n\n - Given that the formulation applies to finite-horizon MDPs, how could this be extended to infinite-horizon settings?\n\n- The occupancy estimation relies on sampling, which can introduce high errors. Have the authors considered methods to reduce estimation errors?\n\n - How does the method perform in more complex tasks beyond the three considered in the experiments? Why were additional MuJoCo tasks not included?\n\n- How does this work compare to prior model-free approaches in imitation learning and inverse reinforcement learning?\n\n - Could the authors discuss the scalability of their approach when applied to larger or more complex domains?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DF3zwmgvr4", "reviewer": "Reviewer_H8vF", "summary": "This paper proposes ZILOT, a MPC-style inference-time trajectory optimization technique that can learn from a single incomplete state-only expert demonstration with pretraining of a dynamic model from an unlabeled state-action dataset. The paper theoretically proves prior method, which is using a goal recognizer with a goal-conditioned policy, will encounter myopic issues and is suboptimal. This paper instead uses MPC, with the Sinkhorn distance between state occupancies of the rollout policy and expert demonstrations as the cost optimization objective. The distance metric between two states are the expected steps from one state to another, which is learned through a goal-conditioned value function. The optimal transport between occupancies are discretized into matching between the states. To determine whether expert states are reachable, another value function is learned to provide an estimate. On several testbeds, the proposed method outperforms existing baselines.\n\n## Update After Rebuttal\n\nI have updated my score accordingly during the author-discussion period; no further update.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Definition of zero-shot\nWe 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.\n\nExisting 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.\n\nThis 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.\n\nWe thank the reviewer for encouraging this discussion, which we hope clarifies our definitions. We will update our submission accordingly.\n\n# Concerns - Finite sample approximation of OT\nWe 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.\n\n# Suggested Improvements\n## World Model Training\nWe 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.\n\n## Goal Abstraction\nWe 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].\nZILOT’s planning objective (Eq. 17) matches the sampled state occupancy to the expert goal occupancy, and thus uses a standard GCVF $V(s, g)$.\n\nEq. 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.\n\n## Extra Evaluation Environments\nPlease refer to our response to Reviewer LRC1, where we motivate our choice of environments and provide evaluations of the Walker environment.\n\n# Other Weaknesses - Running Time\nWe 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].\n\n# Question on Expert mismatch\nAs 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).\n\n---\nWe hope this addresses your questions, and we are happy to elaborate on points that were inadequately covered due to character limitations.\n\nReferences:\n1. Pirotta et. al. Fast Imitation via Behavior Foundation Models. ICLR ‘24\n2. Pathak et. al. Zero Shot Visual Imitation. ICLR ‘18\n3. Touati et. al. Learning one Representation to optimize all Rewards. NeurIPS ‘21\n4. Genevay et. al. Sample Complexity of Sinkhorn Convergences. AISTATS ‘19\n5. Weed et al. Sharp Asymptotic and Finite-Sample Rates of Convergence of Empirical Measures in Wasserstein Distance.\n6. Andrychowicz et. al. Hindsight Experience Replay. NeurIPS ‘17\n7. Cuturi et. al. Optimal Transport Tools (OTT) arXiv:2201.12324\n8. Feydy. et. al. Interpolating between Optimal Transport and MMD using Sinkhorn Divergences. AISTATS ‘19", "claims_and_evidence": "Overall, I feel the claim \"zero-shot\" is somewhat strange to me; I think this is one of my major concern for this paper. In this paper, the definition of \"zero-shot\" seems to be \"(a single) demonstration devoid of actions (which is rough and partial), such as specifying only a few checkpoints\" (Sec. 1). \n\nHowever:\n\n1. in one of the zero-shot prior work the authors referred to in Sec. 1, Pathak et al. [1], the definition of zero-shot is \"the agent never has access to expert actions during training or for the task demonstration at inference\" (as specified in the abstract). This is, however, not the case in this paper; $\\mathcal{D}\\_\\beta$ contains actions as specified in Sec. 2.1. \n\n2. In another paper the authors referred to in Sec. 1, Pirotta et al. [2], the author writes in their Sec. 1: \"1) ... no prior knowledge or demonstrations of the behavior to be imitated are available, and only a dataset of unsupervised trajectories is provided (which is the case for this paper); 2) ... solve any imitation task without any additional samples on top of the demonstrations, and without solving any complex RL problem (which means offline, but not necessarily without actions) ... computation needed to return the imitation policy should be minimal (which is questionable for this work)\".\n\n3. in \"zero-shot\" reinforcement learning [3], which Pirotta et al. [2] refers to (as Pirotta et al. did not formally define the word \"zero-shot\"), the definition of zero-shot is \"solve any RL task in a given environment, instantly with no additional planning or learning, after an initial reward-free learning phase\", which is also not satisifed by this paper.\n\n4. The authors also use LLM's in-context learning ability as an example for zero-shot capabilities in the beginning of this paper. However, LLM's in-context learning ability is \"no any training related to the problem\", which is different from this paper with a world model training stage.\n\n5. according to the author's definition in this paper (state-action offline dataset + state-only incomplete / goal-only trajectories), there are many other papers that fits into this definition in the name of \"learning from observations\", such as IQLearn [4], TAILO [5], PW-DICE [6], SMODICE [7], SAIL [8], RCE [9] (which is goal-based IL), AILOT [10]. The authors seem to only have cited SMODICE in the appendix as an example of forward-backward framework.\n\nTherefore, I feel the word \"zero-shot\" is not well-defined in this paper.\n\n\n**Reference**\n\n[1] D. Pathak et al. Zero-Shot Visual Imitation. In ICLR, 2018.\n\n[2] M. Pirotta et al. Fast Imitation via Behavior Foundation Models. In ICLR, 2024.\n\n[3] A. Touati et al. Does Zero-Shot Reinforcement Learning Exist? In 34th Offline Reinforcement Learning Workshop at Neural Information Processing Systems, 2022.\n\n[4] D. Garg et al. IQ-Learn: Inverse soft-Q Learning for Imitation. In NeurIPS, 2021.\n\n[5] K. Yan et al. A Simple Solution for Offline Imitation from Observations and Examples with Possibly Incomplete Trajectories. In NeurIPS, 2023. \n\n[6] K. Yan et al. Offline Imitation from Observation via Primal Wasserstein State Occupancy Matching. In ICML, 2024.\n\n[7] Y. J. Ma et al. Versatile Offline Imitation from Observations and Examples via Regularized State-Occupancy Matching. In ICML, 2022.\n\n[8] F. Liu et al. State Alignment-based Imitation Learning. In ICLR, 2020.\n\n[9] B. Eysenbach et al. Replacing Rewards with Examples: Example-Based Policy Search via Recursive Classification. In NeurIPS, 2021.\n\n[10] M. Borbin et al. Align Your Intents: Offline Imitation Learning via Optimal Transport. In ICLR, 2025.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria over make sense for the problem. There are a few possible concerns and improvements though:\n\n**Concerns**\n\nIn the paper, the authors try to estimate the state occupancy using $\\rho^\\pi_N\\approx\\frac{1}{N+1}\\sum\\_{t=0}^N\\delta\\_{s\\_t}$. While empirically, works such as OTR [1] mentioned in the paper have proved the success of the discretization, it is worth noting that the estimation of optimal transport between distributions from discrete samples are very inaccurate [2, 3]; in fact, one will need exponentially many samples with respect to the number of state dimensions to accurately estimate the Wasserstein distance between distribution. Thus, the story of approximation from Eq. 12 to the end of \"occupancy estimation\" does not sound convincing enough.\n\n**Improvements**\n\n1. The paper did not specify how the \"world model\" is trained from the offline dataset, which I feel is an important part of the method (as the authors indicated in Appendix C.4, it seems to follow prior work, but the paper should still introduce the method and its implementation in this paper to be self-contained). Moreover, the paper did not specify what a goal abstraction $\\phi(\\cdot)$ is and how it is implemented in the proposed method or baselines (given this, it is unexplained on the objective right above Eq. 18 why $W(\\phi(s);\\phi(s'))$ has goal abstraction on both state inputs but $V(s,\\phi(s'))$ only has goal abstraction on the second state input. This is another major concern of mine.\n\n2. It will be better if the authors can show the result of ZILOT on more environments (walker2d, ant, antmaze, franka kitchen etc.), which are usually accompanieed with the halfcheetah environment in the literature.\n\n**References**\n\n[1] Y. Luo et al. Optimal Transport for Offline Imitation Learning. arXiv: In ICLR, 2023.\n\n[2] J. Stanczuk et al. Wasserstein GANs Work Because They Fail (to Approximate the Wasserstein Distance). arXiv:2103.01678.\n\n[3] J. Weed et al. Sharp Asymptotic and Finite-Sample Rates of Convergence of Empirical Measures in Wasserstein Distance. arXiv:1707.00087.", "theoretical_claims": "The paper has one theoretical claim in Sec. 3, which is about the suboptimality of the method proposed in Pathak et al. [1]. The authors claim that there exists some controllable Markov chain and sequence of goals that the existing method, even with optimal solution, cannot all achieve. The authors give a simple counterexample as a proof and it seems correct to me. To make it more rigorous, I would suggest add \"almost surely\" or \"with probability 1\" to the conclusion \"will not reach all goals in the sequence\".\n\n**Reference**\n\n[1] D. Pathak et al. Zero-Shot Visual Imitation. In ICLR, 2018.", "experimental_designs_or_analyses": "Yes, I think the experiment are generally solid as many details as well as visualizations are given in the appendix. The results seem reasonable which shows the proposed method outperforms several baselines, and ablations are made in Sec. 5.3. Many environment variants are tested in Tab. 2, which I feel sufficient to prove the performance of the proposed method.", "relation_to_broader_scientific_literature": "This paper is beneficial for the Reinforcement Learning (RL) / Imitation Learning (IL) community and robotics community. It does not have significant impact for researchers out of these communities.", "essential_references_not_discussed": "I suggest the authors check \"learning from observations\" in the imitation learning literature, as mentioned in \"Claims and Evidence\" point 5. Currently, I feel multiple papers are missing from the literature discussion. Also, I would suggest the authors to discuss the relation of this work with the ICVF paper [1], which also uses a goal-conditioned value function and aims to learn dynamics through value function from action-free data, which is similar to the goal-conditioned valuer function in this paper. \n\n**References**\n\n[1] D. Ghosh et al. Reinforcement Learning from Passive Data via Latent Intentions. In ICML, 2023.", "other_strengths_and_weaknesses": "**Other Strength**\n\nOverall, I feel the high-level idea of this paper is clearly conveyed and easy to follow. The shortcoming of the prior methods are clearly illustrated in Fig. 1 and rigorously stated afterwards in Sec. 3.\n\n**Other Weakness**\n\nThe method needs to solve optimization problem for every step, which is inefficient; as the authors suggested in Appendix C.3, the method only runs at 0.5 to 3Hz, which is too slow for practical uses (and even slower for simulations).", "other_comments_or_suggestions": "See the other parts of the review.", "questions_for_authors": "I have two question for the authors: \n\n1. Is the proposed method robust against dynamic mismatch between expert demonstration and the agent? I would expect the method to be robust if it can work with incomplete state-only trajectories and uses Wasserstein distance to match those states. An example of such experiment can be found in Appendix H of the SMODICE [1] paper.\n\n**References**\n\n[1]  Y. J. Ma et al. Versatile Offline Imitation from Observations and Examples via Regularized State-Occupancy Matching. In ICML, 2022.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper studies the setting of zero-shot imitation learning, where an imitation policy is directly computed from a single demonstration. The authors shows how translating this problem into a "greedy" goal-based problem may return poorly-performing policies and suggest ways to mitigate the issue by leveraging distance in occupancy distributions (approximated by a world model) to drive the imitation learning objective.

There is general consensus among the reviewers of the interest of the theoretical analysis revealing the potential limitations of the goal-based reduction of imitation learning. After the rebuttal, the authors have partially addressed concerns on the empirical validation by complementing the results initially available in the paper and by clarifying aspects related to the planning process. Based on this, I recommend acceptance for the paper, but I strongly encourage the authors to incorporate the comments of Reviewer LRC1. In particular, I believe that further ablations on OT vs f-divergences can be easily addressed. Furthermore, I tend to agree with the authors that a good variety of domains is already covered in the current submission, I still encourage them to assess whether expanding the empirical validation to other domains is possible.
Decision: Accept (poster)