Fast Imitation via Behavior Foundation Models
Authors: Matteo Pirotta
Andrea Tirinzoni
Ahmed Touati
Alessandro Lazaric
Yann Ollivier
Venue: ICLR 2024 spotlight (2024)
Keywords: Behavior Foundation Models
unsupervised reinforcement learning
imitation learning
Primary area: reinforcement learning

Abstract: Imitation learning (IL) aims at producing agents that can imitate any behavior given a few expert demonstrations. Yet existing approaches require many demonstrations and/or running (online or offline) reinforcement learning (RL) algorithms for each new imitation task. Here we show that recent RL foundation models based on successor measures can imitate any expert behavior almost instantly with just a few demonstrations and no need for RL or fine-tuning, while accommodating several IL principles (behavioral cloning, feature matching, reward-based, and goal-based reductions). In our experiments, imitation via RL foundation models matches, and often surpasses, the performance of SOTA offline IL algorithms, and produces imitation policies from new demonstrations within seconds instead of hours.

Published as a conference paper at ICLR 2024
FAST IMITATION VIA BEHAVIOR FOUNDATION MODELS
Matteo Pirotta", Andrea Tirinzoni' & Ahmed Touati"
Alessandro Lazarie & Yann Ollivier Fundamental AT Research Meta { (lazaric, yol @meta com
Fundamental AT Research at Meta (pirotta, t: tirinzoni, atouati @meta. com
ABSTRACT
Imitation learning (IL) aims at producing agents that can imitate any behavior given few expert demonstrations. Yet existing approaches require many demonstrations and/or running (online or offline) reinforcement learning (RL) algorithms for each new imitation task. Here we show that recent RL. foundation models based on successor measures can imitate any expert behavior almost instantly with just few demonstrations and no need for RL or fine-tuning, while accommodating several IL principles (behavioral cloning, feature matching, reward-based, and goal-based reductions). our experiments, imitation via RL foundation models matches, and often surpasses, the performance of SOTA offline IL algorithms, and produces imitation policies from new demonstrations within seconds instead of hours.
INTRODUCTION
The objective of imitation learning (Schaal, 1996, II.) is develop agents that can imitate any behavior from few demonstrations. For instance, a cooking robot may learn how to prepare a new recipe from single demonstration provided by an expert chef. A virtual character may learn to play different sports in virtual environment from just few videos of athletes performing the real sports. Imitation learning algorithms achieved impressing results in challenging domains such as autonomous car driving (Bühler et al., 2020; Zhou et al., 2020; George et al., 2018), complex robotic tasks (Nair et al., 2017: Lioutikov et al., 2017: Zhang et al., 2018: Peng et al., 2020; Mandi al., 2022; Pertsch et al., 2022; Haldar et al., 2023), navigation tasks (Hussein et al., 2018; 8; Shou et al., 2020), cache management (Liu al., 2020), and virtual character animation (Zhang et al., 2023; Peng al., 2018; Wagener al., 2023). Despite these achievements, existing approaches (see Sec. 2 for detailed review) suffer from several limitations: for any new behavior to imitate, they often require several demonstrations, extensive interaction with the environment, running complex reinforcement learning routines, or knowing in advance the family of behaviors to be imitated.
In this paper, we tackle these limitations by leveraging behavior foundation models (BFMs)' to accurately solve imitation learning tasks from few demonstrations. To achieve this objective, we BFM to have the following properties: 1) When pre-training the BFM, no prior knowledge or demonstrations of the behaviors to be imitated are available, and only a dataset of unsupervised transitions/trajecooiiss is provided; 2) The BFM should accurately solve any imitation task without any additional samples on top of the demonstrations, and without solving any complex reinforcement learning (RL) problem. This means that the computation needed to return the imitation policy (i.e., the inference time) should be minimal; 3) Since many different ways to formalize the imitation learning problem have been proposed (e.g., behavior cloning, apprenticeship learning, waypoint imitation), we also want a BFM that is compatible with different imitation learning settings.
Our main contributions can be summarized as follows.
We leverage recent advances in BFMs based on successor measures, notably the forward-backwant (FB) framework (Touati et al., 2023; Touati & Ollivier, 2021), to build BFMs that can be used solve any imitation task, and satisfy the three properties above. We focus on FB for its
"Joint first author, alphahetical order.
Joint last author, alphahetical order.
'The term "Behavior" emphasizes that the model aims at controlling agent in dynamical environment. This avoids confusion with widely used foundation models for images, videos, motions, and language. See Yang al. (2023) for extensive review of the latter for decision making.

---

Published as a conference paper at ICLR 2024
BPM
FB-L (cun)
700
26.2
Imtattnn Score (5)
30.9
406
128
8C
ORd
OTR
SQn
TEaHL
DEMDDICE
IQLEaRM
Br
ER BRoaan
GoMtetu
COAL PPg
Mannry
GoaL-CP CPT
BC
REita
Bii Srt
FM
ORI
1Q
Figure 1: Imitation score (ratio between the cumulative return of the algorithm and the cumulative return of the expert) averaged over domains, tasks, and repetitions, for a single expert demonstration. FB-IL methods reach SOTA performance with fraction of the test-time computation needed by offline IL. baselines and they perform better than other pre- pre-trained behavior foundation models, while implementing wider set of principles. Notice that goal-based methods, such as GOAL-TD3 and GOALFB, work well in this experiment but have more restricted applicability (Sec. 4.4).
demonstrated performance zero-shot reinforcement learning compared to other approaches (Touati et al., 2023). We refer to the set of resulting algorithms as FB-IL.
We test FB-IL algorithms across environments from the DeepMind Control Suite (Tassa et al., 2018a) with multiple imitation tasks, using different IL principles and settings. We show that not only do FB-IL algorithms perform on-par or better than the corresponding state-of-the-art offline imitation learning baselines (Fig. 1), they also solve imitation tasks within few seconds, which three orders of magnitude faster than offline IL methods that need to run full RL routines compute an imitation policy (Fig. 2). Furthermore, FB-IL IL methods perform better than other BFM methods, while being able to implement a much wider range of imitation principles.
RELATED WORK
While a thorough literature review and classification is out of the scope of this work, we recall some of the most popular formulations of IL. each of which will be implemented via BFMs Sect. 4.
Offline IL
Time
BC
3h14m
TD3-IL
7h3m
Behavioral Cloning (Pomerleau, 1988; Bain & Sammut, 1995, BC) aims at directly reproducing the expert policy by maximizing the likelihood of the expert actions under the trained imitation policy. While this is the simplest approach to IL, it needs access to expert actions, it may suffer from compounding enrors caused by covariate shift (Ross et al., 201 1), and it often requires many demonstrations learn an accurate imitation policy. Variants of the formulation include regularized BC (c.g., Piot et al., 2014; Xu et al., 2022) and BC from observations only (e.g.. Torabi et al., 2018).
Demodice
12h59m
FB-IL
Time
BCFB
Im
ERrB
BBELLFB
4m
Figure 2: Time for computing an imitation policy from sin- gle demonstration for subset of offline IL basclines and FB- IL methods, averaged over all environments and tasks.
Another, simple but effective, IL principle is to design (e.g., Ciosek, 2022: Reddy et al., 2020) or infer a reward (c.g.. Zolna et al., 2020; Luo et al., 2023; Kostrikov et al., 2020) from the demonstrations, then use it to train RL. agent. For instance, SQIL (Reddy et al.,
2020) assigns reward of to expert samples and 0 to non-expert samples obtained either from an offline dataset or online from the environment. Other methods (e.g., Ho Ermon, 2016; Zolna et al., 2020; Kostrikov et al., 2020; Kim et al., 2022b;a; Ma et al., 2022) learn discriminator to infer reward separating expert from non non-expert samples. OTR (Luo et al., 2023) uses optimal transport to compute distance between expert and non-expert transitions that is used reward. Finally, other approaches frame IL as a goal-condine condicondt task (e.g.. Ding et al., 2019; Lee et al., 2021), and leverage advances in gral-oriented RL by using goals extracted from expert trajectories.
2

---

Published as a conference paper at ICLR 2024
Many imitation learning algorithms can be formally derived through the lens of either Apprenticeship Learning (AL) (e.g., Abbeel & Ng. 2004; Syed & Schapire, 2007; Ziebart et al., 2008; Syed et al., 2008; Ho et al., 2016; Ho & Ermon, 2016; Garg et al., 2021; Shani et al., 2022; Viano et al., 2022; Al-Hafez et al., 2023; Sikchi et al., 2023) or Distribution Matching (DM) (e.g., Kostrikov et al., 2020; Kim et al., 2022a; Zhu et al., 2020; Kim al., 2022b; Ma et al., 2022; Yu et al., 2023). AL looks for a policy that matches or outperforms the expert for any possible reward function in a known class. If the reward is linearly representable w.r.f. set of features, a sufficient condition is to find a policy whose successor features match that of the expert. DM approaches directly aim to minimize some f-divergence between the stationary distribution of the learned policy and the one of expert.
The main limitation of these approaches is that they need to solve a new (online or offline) RL problem for each imitation task from scratch. This often makes their sample and computational complexity prohibitive. Brandfonbrener et al. (2023) pre-train inverse dynamics representations from multitask demonstrations, that can be efficiently fine-tuned with BC to solve some II. tasks with reduced sample complexity. Masked trajectory models (Carroll et al., 2022; Liu et al., 2022; Wu et al., 2023) pre-train transformer-basee architectures using random masking of trajectories and can perform
imitation if provided with sufficiently curated expert datasets at pre-training In a similar setting, Reuss et al. (2023) use pre-trained coal-conditioned policies based on diffusion models for waypoint-conditionee imitation. Wagener al. (2023) pre-train autoregressive architectures with multiple experts, but focuses on trajectory completion rather than "full" imitation learning. One-shot IL (e.g.. Duan al.. 2017: Finn al., 2017; Yu et al., 2018; Zhao et al., 2022; Chang & Gupta, 2023) uses meta-learning to provide fast adaptation to new demonstration train time. This requires carefully curated datasets train time with access to several expert demonstrations. Task- consitondt approaches (e.g., James et al., 2018; Kobayashi et al.. 2019; Dasari & Gupta, 2020; Dance et al., 2021) can solve IL by (meta- learming model conditioned on reward, task or expert embeddings by accessing privileged information at train time.
coyypoint-condiiionn
3 PRELIMINARIES
Markoy decision processes. Let M = (S, P. be a reward-free Markov decision process (MDP), where S is the state space, A is the state space, P(ds'|s, a) is the probability measure on 8' E S defining the stochastic transition to the next state obtained by taking action a in state 8, and 0 7 < is discount factor (Sutton & Barto, 2018). Given (sp,
S x A and a policy n: S + Prob(A), we denote Pr(-Iso, ao, T) and B[-|so, 40. T] the probabilities and ex- pectations under state-action sequences (st, az) starting at (5o, ag) and following policy æ in the environment, defined by sampling pa (ds'ls) := P(ds')s |s, (da
~ P(ds a:11) and ar ~ ((dag(s)). We define (s), the state transition probabilities induced by л. Given a reward function r: S + E. the Q-function of æ for r is (7 (o) := =Ez0 "Err(s)) (so. 4g x]. The optimal Q-function is (5(5,a) a) == sup, Q=(s,a). (For simplicity, depends on $c+1 instead on the full triplet (st, ae, St+1), but this is not essential.)
assume the reward only
For each policy and each so E S, ao E A, the successor measure = dp. over S describes the cumulated discounted time spent at each state 8141 if starting at (So, au) and following n, namely,
M (Sg, ag. X) = Σ12 ? Pr(se+1 E X (0 ao T) 4X C.
(1)
The forward-backward (FB) framework. The FB framework (Touati & Ollivier, 2021) learns tractable representation of successor measures that provides approximate optimal policies for any reward. Let Rd be a representation space, and let p be an arbitrary distribution over states, typically the distribution of states in the training set. FB learns two maps F: S Ax Rd B: S + Idd and a set of parametrized policies (TT)ZERd, such that
Rd and
M 0n X) [x F(so. og. =)TT B(s) p(ds), (s) R arg max, F(s,a,z)"=
Vso E S, ag E A, C S, E Rd as E A, = Rdd.
We recall some propertics of FB that will be leveraged to derive FB-based imitation methods. In the following, use the short forms Cov B = (B)) and
Md (s) == Ba (s) Md (s)
F F(s, Bo)x((s) a(s,a,2).
(3)
AA few papers (e.g., Peng al., 2022; Juravsky al., 2023) have used expert trajectories to speed up the learning of task-conditioned policies. Their objective is not II. but task generalization and/or compositionality
3

---

Published as a conference paper at ICLR 2024
Proposition 1 (Touati & Ollivier (2021)). Assume (2) holds exactly. Then the following holds.
First, for any function r: S R. let
24 = Bexplr(s)B(ss)
(4)
Then 2,, is optimal for i.e., T2, E arg max, Q7 (8, a). Moreover,
Finally, for each policy и; and each (so. 0n) E S A, ((su, ao. 2) E Rd are the successor features associated to the state embedding ((()) = (Cov B(s), i.e.
F(so, ao; 2) EE Ei20 f (so, ao, T2
(S)
In practice, the properties in Prop. I only hold approximately, as F' B is a rank-d model of the successor measures, n. may not be the exact greedy policy, and all of them are learned from samples. Moreover, (4) expresses as an expectation over states from the training distribution If sampling from a different distribution p' at test time, an approximate formula is (Touati & Ollivier, 2021, $B.5):
= (E (()) (T B(s) E [())
(6)
Pre-training an FB model can be done from a non-curated, offline dataset of trajectories or transitions, thus fulfilling property 1) above. Training is done via the measure-valued Bellman equation satisfied by successor measures. We refer to (Touati et al., 2023) for a full description of FB training.
FB belongs to wider class of methods based on successor features (e.g.. Borsa et al. (2018)). Many of our imitation algorithms still make sense with other methods in this class, see App. A.6. We focus on FB as it has demonstrated better performance for zero-shot reinforcement learning within this family (Touati et al., 2023).
4 FORWARD-BACKWARD METHODS FOR IMITATION LEARNING
We consider the standard imitation learning problem, where we have access to a few expert trajectories (So, ......, S2(r)), each of length ((T), generated by some unknown expert policy Te. and no reward function available. In general, we do not need access to the expert actions, except for behavioral cloning. We denote by E, the empirical average over the expert trajectories and by Pe the empirical distribution of states visited by the expert trajectories."
We describe several IL methods based a pre-trained FB model. These only from demonstration data, without solving any complex RL problem at test time (property 2)). Some methods just require near-instantaneous forward pass through B test time, while others require parameter z. The latter is still much faster than solving
gradient descent over the small-dimensione full RL problem, as shown in Fig. 2. At imitation time, we assume access the functions F, B, the matrix Cov B, and the policies т., but we do not reuse the unsupervised dataset used for FB training. To illustrate how FB can accommodate different principles, we present the methods in loose groups by the underlying IL principle.
4.1 BEHAVIORAL CLONING
In case actions are available in the expert trajectories, we can directly implement the behav- ioral cloning principle using the policies (Ta); returned by the FB model. Each policy Ra defines a probability distribution on state-action sequences given the initial state Sp, namely Pr(au, $1, d1,... ! π.) = II T2 Tz(a, P(s a.). We look for the n, for which the expert trajectories most likely, by minimizing the loss
LBC(z) == - E, In Pr((ag, $1, 21.41.... So. nI) = - Σ In T. (aglst) + est,
where the constant absorbs the environment transition probabilities P(ds (s,, at), which do not depend on 2. Since we have access to π. (als), this can be optimized over 2 given the expert trajectories, leading to the behavior cloning- FB (BCrB) approach.
*e give each expert trajectory the same weight in pe independently of its length, pe corresponds to first sampling trajectory, then sampling state in that trajectory.
4

---

Published as a conference paper at ICLR 2024
Since the FB policies (Tz)); are trained to be approximately optimal for some reward, expect FB (and BFMs in general) to provide a convenient "bias" to identify policies, instead of performing BC among the set of all (optimal or not) policies.
4.2 REWARD-BASED IMITATION LEARNING
Existing reward-based IL methods require running RL algorithms to optimize an imitation policy based on reward function specifically built to mimic the expert's 's behavior. Leveraging FB models, we can avoid solving an RL problem at test time, and directly obtain the imitation policy via simple forward pass of the B model. Indeed, as mentioned in Sec. FB models can recover a (near-optimal) policy for any reward function r by setting z Expp (r(s) B()) Depending on the specific reward function, we obtain the following algorithms to estimate 2, after which we just use ne.
First, consider the case of r-) pe() in (Kim et al., 2022b;a; Ma et al., 2022). This yields
z Ex lr (()) E,, [ E,
[( Cr20 ((s--1)
(8)
which amounts to using the FB formula (4) for by just putting a reward at every state visited by the expert. We refer to this empirical reward via FB (ERyB).
Similarly, the reward r(-) = p(t)) + pc(-)) used in (Reddy et al., 2020; Zolna et al., 2020) leads to regularized empirical reward via FB (RERFB), derived from (6) in App. A.I:
2 Cov Cov( + Espp, [B(s) (B(s) [B]
Even though these reward functions are defined via the distribution p of the unsupervised dataset, this can be instantiated using only the pre-trained FB model, with no access to the unsupervised dataset, and no need to train discriminator.
Note that (8) and (9) are independent of the order of states in the expert trajectory. This was not problem in our setup, because the states themselves carry dynamical information (speed variables). If this proves limiting in some environment, this can easily be circumvented by training successor measures over visited transitions (8t, 8(+1) rather than just states 8t+1, namely, training the FB model with (s,, St+1). A similar trick is applied, e.g., in (Zhu et al., 2020; Kim et al., 2022a).
4.3 DISTRIBUTION MATCHING AND FEATURE MATCHING
Apprenticeship learning and distribution matching are popular ways to provide a formal definition of II. as the problem of imitating the expert's visited states. We take unified perspective on these two categories and derive several FB-IL methods starting from the saddlc-point formulation of IL common to many AL and DM methods. Let po be arbitrary initial distribution over S. For any reward r and policy т, the expected discounted cumulated return of π is equal to Ex pa (M))) by definition of M Consequently, the AL criterion of minimizing the worst-case performance gap between π and the expert can be seen as a measure of divergence between successor measures:
inf sup Ex <M (sa), ) r) (M* (so), r)] = inf [Essoppo =**((s) Esorpo M (3o) R. (10) rER where R any class of reward functions, and ||- the resulting dual seminorm. Since FB directly models M7, it can directly tackle (10) as finding the policy Tu that minimizes the loss
Ex- (z) == [[Eeo<po = (so) Beo M (so))l
(11)
In practice, instead of (11), we consider the loss
Cx- (z) == Eeorrp [IM - M (so))
(12)
This is stricter criterion than (11) as it requires the successor measure of the imitation policy and the expert policy to be similar for any s observed along expert trajectories. This avoids undesirable effects from averaging successor measures over Po. which may "erase" too much information about the policy (e.g., take S = [$1,32] where one policy swaps s and 82 and the other policy does nothing: average over the starting point, the two policies have the same occupation measure). This increases robustness in our experiments (see App. E.6).
We can derive wide range of algorithms depending the choice of R, how we estimate M from expert demonstrations, and how we leverage FB models to estimate = For instance, our algorithms can be extended to the KI. divergence between the distributions (App. A.5).
5

---

Published as a conference paper at ICLR 2024
Successor feature matching. A popular choice for R is to consider rewards linear in given feature basis (Abbeel & Ng, 2004). Here we can leverage the FB property of estimating optimal policies for rewards in the linear span of B (Touati et al., 2023). Taking RB {r w B, R | ||2 S 1} in (10) yields the seminorm || = SuPrER J r(s))((ss) ||. (ds) and the loss
CB (2) = Eso-pe IIS B(5) (50, (50, - B(s) - (so. de)
(13)
namely, the averaged features B of states visited under nı and ne should match. This can be computed by using the FB model for M and the expert trajectories for M as follows.
Theorem 2. Assume that the FB successor feature property (5) holds. Then the loss (13) satisfies
LB (2) = Esprp. E [[(CoB)F(ss Exzo x B(SUHE+)) Se, He
+ est.
(14)
This can be estimated by sampling a segment (st, 8141,...) starting at a random time t on an expert trajectory. Then we can perform gradient descent over z. We refer this method as FMFB.
Distribution matching. If R is restricted to the span of some features, we only get seminorm on successor measures (any information not in the features is lost). Instead, one can take = (p). which provides a full norm ||M M (p)) on visited state distributions: this matches state
distributions instead of features. This can be instantiated with FB (App. A.3), but the final loss is very similar to (67), as FB neglects features outside of anyway. We refer to this method as DMFn-
Bellman residual minimization for distribution matching. An alternative approach is identify the best imitation policy or its stationary distribution via the Bellman equations they satisfy. This is to distribution matching what TD is direct Monte Carlo estimation of funcction..
The successor measure of a policy satisfies the measure-valued Bellman equation M M ds') = pa (st, ds') + S pa P" (st, de 1) (s++1,,,,, or more compactly M pa +7 M (Blier et al., 2021). So the successor measure of the expert policy satisfies M = p + 7m M* Therefore, if we want to find policy т. that behaves like Te, = should approximately satisfy the Bellman equation for p",, namely, nM p + p M Thus, we can look for policy x; whose Bellman gaps for Tee small. This leads to the loss
Be)Bel(2)) == ||M - p - yP M=
(15)
where the bar above M on the right denotes stop-grad operator, usual for deep Q-learning.
The method we call BBELLPB uses the seminorm |1-1 in (15). This amounts to minimizing Bellman gaps of Q Q-fuctioons for all rewards linearly spanned by B. With the FB model, the loss (15) with this norm takes a tractable form allowing for gradient descent over z (App., Thm. 4):
Ln))el())) = =E
s1) [- (s,, =) (Cov B)B(s++1
+ (F(st, 7 F( (Cov B)² -FF + est. (16)
The norm from R= (p) in (15) yields loss similar to the one used during FB training (indeed, FB is trained via a similar Bellman equation with næ instead of т.). The final loss only differs from (16) by Cov B factors, so we report it in App. A.4 (Thm. 5). We call this method FBLOSSYB.
Relationship between IL principles: loss bounds. Any method that provides a policy close to Te will provide state distributions close that of π. as a result, so we expect a relationship between the losses from different approaches. Indeed, the Bellman gap loss bounds the distribution matching loss (12), and the BC loss bounds the KI. version of (12). This is formalized in Thms. 7 and (App. A.7).
4.4 IMITATING NON-STATIONARY BEHAVIORS: GOAL-BASED IMITATION
While most IL methods designed imitate stationary behaviors, we can leverage FB models to imitate non-stationary behaviors. Consider the case where only a single expert demonstration т is available. At each time step I, use the FB method reach state 8144 slightly ahead of

---

Published as a conference paper at ICLR 2024
St in the expert trajectory, where & 2 0 is a small, fixed integer. Namely, we place single reward at $1+k; use the FB formula (4) to obtain the 24 corresponding to this reward, 2, : (( ) and use the policy z.. We call this method GOALFB- This is related to settings such tracking (e.g.. Wagener et al., 2023; Winkler et al., 2022), waypoint imitation (e.g.. Carroll et al., 2022; Chang & Gupta, 2023; Shi et al., 2023), goal-based II. (e.g., Liu et al., 2022; Reuss al., 2023).
GOALEB leverages the possibility to change the reward in real time with FB. A clear advantage is its ability to reproduce behaviors that do not correspond to optimizing a (Markovian) reward function, such as cycling over some states, or non-stationary behaviors. GOALER may have advantages even in the stationary case, as it may mitigate upproximation errors from the policies or the representation of M": by selecting a time-varying 2, the policy can adapt over time and avoid deviations in the execution of long behaviors through stationary policy. However, goal-based II. limited to copying one single expert trajectory, by reproducing the same state sequence. The behavior cannot necessarily be extended past the end of the expert trajectory, and no reusable policy extracted.
5 EXPERIMENTS
In this section, we evaluate FB-IL against the objectives stated in the introduction: Property 2. We verify if an FB model pre-trained on one specific environment is able to imitate wide range of tasks with only access to few demonstrations and without solving any RL problem. Property 3. We assess the generality of FB-IL by considering a variety of imitation learning principles and settings.
Protocol and baselines. We evaluate IL methods on 21 tasks in 4 domains (Maze, Walker, Cheetah, Quadruped) from (Touati et al., 2023). We use the standard reward- based evaluation protocol for IL. For each task, function (Tassa et al., 2018a). We use the expert policies to generate 200 trajectories for each task to be used for II.. In our first series of experiments, the II. algorithms demonstration (see App. R.4 for the effect of additional demonstrations). Each experiment (i.e., pair algorithm and task) was repeated with 20 random seeds. We report the cumulated reward achieved by the IL policy, computed using the ground-truth task-specific reward and averaged over 1000 episodes starting from the same initial distribution used to collect the expert demonstrations.
train expert policies using TD3 (Fujimoto et al., 2018)
task-specific reward
provided with single expert
For each environment, we train an FB model using only unsupervised samples generated using RND (Burda et al., 2019). We repeat the FB pre-training 10 times, and report performance averaged over the resulting models (variance is reported in App. E.3). For FB-II. methods that require gradient descent over = (BCrB, BBEI.I-B, and FMFB), we use warm-start with 2g = E ( (Te)) )) (8), which can be computed with forward passes on B only. GOALEB is run with a lookahead window h = 10.
First (Section 5.1), we compare FB-IL to standard offline IL algorithms trained on each specific imitation task, using the same unsupervised and expert samples as FB-IL. For behavioral cloning approaches, we use vanilla BC. For reward-based IL, we include SQIL (Reddy et al., 2020) (which is originally online but can easily be adapted offline; SQIL balances sampling in the update step and runs SAC): TD3-IL -I (where we merge all samples in the replay buffer and use TD3 instead of SAC); ORIL (Zolna al., 2020) from state-action demonstrations and only state; and OTR (Luo et al., 2023) using TD3 as the offline RL subroutine. For AL and DM II., we use DEMODICE (Kim et al., 2022b), and IQLEARN (Garg et al., 2021). See App. B for details.
Next (Section 5.2), we also include alternative behavior foundation models beyond FB, pre-trained for each environment on the same unsupervised samples FB. GOAL-TD3 pre-trains goal- conditioned policies π (a]ss,) g) on the unsupervised dataset using TD3 with Hindsight Experience Replay (Andrychowicz et al., 2017). At test time, it can implement goal-based IL, i.c., at each time step t it selects the policy 7 <(a $t,, s(4)) where the goal si corresponds to a state i steps ahead in the expert trajectory. (Despite the simplicity, we did not find this algorithm proposed in the literature.) Next, GOAL-GPT (Liu et al., 2022) pre-trains goal-conditioned, transformer- based auto-regressive policy r(a,(s g), (s:-1,9) g). (+t g); g = St+*) to predict the next action based on last states and the state k steps in the future as the goal of the policy. MASKDP (Liu et al., 2022) uses bidirectional transformer to reconstruct trajectories with randomly masked states and actions. Both models can be used to perform goal- based IL. We adapt DIAYN (Eysenbach et al., 2018) pre-train a set of policies (π₂) with 2 e IR and a skill decoder ; S + Rd predicting which policy is more likely reach a specific state. (This requires online interaction during pre-training.) It can be used to
7

---

Published as a conference paper at ICLR 2024
MR DUL
K Ho Juee)
614
613
N11 AL
312
DMC Ouadiaged
DMC Walker
1
124
BC
Be
BC
ORE
ORII-O
OTEE
2patt
ER
ER
RERe
Pafa
Ooanse Gont
Mr PoaL.PP
Coat Gredes
Bi
Metza
Ellon
FRo
RERo
Fils
Gos Gosteson
Gane Cocce
nu
God Goste
GouL-CP -PP
Figure 3: Imitation score (i.e., ratio between the cumulative return of the algorithm and the one of the
expert) for each domain averaged over tasks and repetitions for a single expert demonstration.
implement behavioral cloning as in (7), method similar ERFB, and goal-based IL by selecting 2, ((((++)). See App. C for extra details.
Additional results. App. E.1 contains detailed results with additional baselines and FB-IL variants. App. E.2 ablates over our warm-start strategy for optimization based FB-IL methods. App. E.4 studies the influence of the number of expert trajectories, with FB methods being the least sensitive, and BC methods the most. App E.5 tests the methods under a shift between the distribution of the initial states for imitation at test time and the one of expert trajectories: the overall picture is largely unchanged from Fig. 1, although the slight lead of goal-based methods disappears. App. E.3 shows that performance is not very sensitive to variations of the pretrained FB foundation model (estimated across 10 random seeds for FB training), thus confirming robustness of the overall approach.
5.1 COMPARISON TO OFFLINE IL BASELINES
Fig. 3 compares the performance of FB-IL methods and offline baselines grouped by IL principle. For ease of presentation, report the performance averaged over tasks of each environment. Overall, FB- IL methods perform on on-par or better than each of the baselines implementing the same IL principle, consistently across domains and IL principle. In addition, FB-IL is able recover the imitation policy in few seconds, almost three orders of magnitude faster than the baselines, that need to be re-trained for each expert demonstration (Tab. 2). This confirms that FB models are effective BFMs for solving a wide range of imitation learning tasks with few demonstrations and minimal compute.
As expected, BC baselines perform poorly with only one expert trajectory. BCFB has a much stronger performance, confirming that the set (m=) contains good imitation policies for large majority of tasks and that they can be recovered by behavioral cloning from even a single demonstration.
Reward-based FB-IL methods -ERFB (8), RERFB (9)- achieve consistent performance across all environments and perform on par or even better than the baselines sharing the same implicit reward function. This shows that FB models are effective at recovering near-optimal policies from rewards. On the other hand, reward-bart based IL offline baselines display a significant variance in their performance across environment (e.g., ORIL for state-action completely fails in maze tasks). The baselines derived from distribution matching and apprenticeship learning perform poorly in almost all the domains and tasks. This may be because they implement conservative offline RI. algorithms that
8

---

Published as a conference paper at ICLR 2024
TW
Method
Score
GOALEB
49.53 + 4.06
GOAL-TD3
50.76 + 6.52
GOAL-GPT
17.58 0.82
Figure 4: Imitating a sequence of yoga poses in the walker domain. (left) Example of sequence. (right) Total reward averaged over 1000 aundomly-generate sequences and 10 pretrained BFMs plus/minus 95% confidence interval. The reward of trajectory is computed by summing the (normalized) rewards of the poses expected at each time step. See App. for additional details.
are strongly biased towards the unsupervised data and fail at imitating the expert demonstrations: + On the other hand, FB-IL variants achieve good results (about 4 times the performance of DICE) in all domains except maze where they lag behind other FB-IL methods. In general, this shows that FB models effective in implementing different IL principles even when offline baselines struggle.
5.2 COMPARISON TO OTHER BFM METHODS
The BFM methods reported in Fig. 3 display a trade-off between generality and performance. DIAYN pre-trained policies and discriminator be used implement wide range of imitation learning principles (except for distribution matching), but its performance does not match the corresponding top offline baselines and it is worse than FB-IL across all domains. Methods based on masked- trajectory models can implement a goal-based reduction of IL and work better than DIAYN. Finally, GOAL-TD3 is performing best among the other BFM methods and it close second w.r.t. GOALFB- Nonetheless, discussed in Sect. 4.4, all goal-based reduction methods are more limited in their applicability, since they can only use single expert demonstration, cannot generalize beyond the expert trajectory, and do not produce policy reusable over the whole space.
5.3 WAYPOINT IMITATION LEARNING
We consider non-realizable and non-stationary experts by generating demonstrations as the concate- nation of "yoga poses" from (Mendonca et al., 2021), implicitly assuming that the expert policy can instantaneously switch between any two poses. We keep each pose fixed for 100 steps and generate trajectories of 1000 steps. In this case, no imitation policy can perfectly reproduce the sequence of poses and only goal-based II. algorithms can be applied, since all other IL methods assume stationary expert policies. We evalute the same pre-trained models used in the previous section.
Fig. 4 shows that GOALFB matches the performance of GOAL-TDE and outperforms GOAL-GPT, This confirms that even in this specific case, FB-IL competitive with other BFM models that specialized (and limited) to goal-reaching tasks, whereas the same pre-trained FB model can be used implement a wide range of imitation learning principles. GP''s poor performance may be because the algorithm tries to reproduce trajectories in the training dataset rather than learning the optimal way to reach goals. We refer to App. for a qualitative evaluation of the imitating behaviors.
CONCLUSION
Behavior foundation models offer a new alternative for imitation learning, reducing by orders of magnitude the time needed to produce an imitation policy from new task demonstrations. This comes the cost of pretraining environment-specffic (but task-agnostic) foundation model. BFMs be used concurenntyy with a number of imitation learning design principles, and reach state-of-the-an performance when evaluated for the ground-truth task reward. One thcoretical limitation that, due to imperfections in the underlying BFM, one may not recover optimal performance even with infinite expert demonstrations. This can be mitigated by increasing the BFM capacity, by improving the training data, or by fine-tuning the BFM at test-time, which we leave to future work.
'In App. G we confirm this intuition by showing that the baselines in this category achieve much better performance when the unsupervised dataset contains expert samples (e.g., D4RL. data). Unfortunately, this requires curating the dataset for each expert and it would not allow solving multiple tasks in the same environment.

---

Published as a conference paper at ICLR 2024
REFERENCES
Pieter Abbeel and Andrew Y. Ng. Apprenticeship learning via inverse reinforcement learning. In ICML, volume 69 of ACM International Conference Proceeding Series. ACM, 2004.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C. Courville, and Mare G. Bellemare. Deep reinforcement learning at the edge of the statistical precipice. In NeurIPS, pp. 29304-29320, 2021.
Firas Al-Hafez, Davide Tateo, Oleg Arenz, Guoping Zhao, and Jan Peters. L.S-IQ: implicit reward regularization for inverse reinforcement learning. In ICLR. OpenReview.net, 2023.
Marcin Andrychowicz, Dwight Crow, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba. Hindsight experience replay. In NIPS, pp. 5058-0045555 2017.
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. abs/1607.064500, 2016.
Layer normalization.
CoRR,
Michael Bain and Claude Sammut. framework for behavioural cloning. In Machine Intelligence 15, pp. 103-129. Oxford University Press, 1995.
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07123, 2021.
Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado van Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. arXiv preprint /rXiv:/812.07022 2018.
David Brandfoombenner, Ofir Nachum, and Joan Bruna. Inverse dynamics pretraining learns good representations for multitask imitation, 2023.
Andreas Bühler, Adrien Gaidon, Andrei Cramariue, Rares Ambrus, Guy Rosman, and Wolfram Burgard. Driving through ghosts: Behavioral cloning with false positives. pp. 5431-5437, 10 2020. 2T0.S444555.222.9944000
doi: 10.1
Yuri Burda, Harrison Edwards, Amos J. Storkey, and Oleg Klimmo. Exploration by random network distillation, In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL https: /oppnr/eee//eeeeeeeeee net/forume ? id 1d=H110JnRR Ym.
Micah Carroll, Orr Paradise, Jessy Lin, Raluca Georgescu, Mingfei Sun, David Bignell, Stephanic Milani, Katja Hofmann, Matthew J. Hausknecht, Anca D. Dragan, and Sam Devlin, Uni[mask]: Unified inference in sequential decision problems. In NeurlPS, 2022.
Matthew Chang and Saurabh Gupta. One-shot visual imitation via attributed waypoints and demon- stration augmentation. In ICRA, pp. 5055-5062. IEEE, 2023.
Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu. Varia- tional empowerment representation learning for goal-conditioned reinforcement learning. In International Conference on Machine Learning, pp. 1953-1963, PMLR, 2021.
Kamil Ciosek. Imitation learning by reinforcement learning. In ICLR. OpenReview.net, 2022.
Christopher R. Dance, Julien Perez, and Théo Cachet. Demonstration Demoontraaion-connittonee reinforcement learning for few-shot imitation. In ICML, volume 139 of Proceedings of Machine Learning Research, pp. 2376-2387. PMLR, 2021.
Sudeep Dasari and Abhinav Gupta. Transformers for one-shot visual imitation. In CoRL, volume 155 Proceedings of Machine Learning Research, pp. 2071-2084, PMLR, 2020.
Yiming Ding, Carlos Florensa, Pieter Abbeel, and Mariano Phiclipp. Goal-conditioned imitation learning. In NewrlPS, pp. 15298-15300 2019.
10

---

Published as a conference paper at ICLR 2024
Yan Duan, Marcin Andrychowicz. Bradly C. Stadie, Jonathan Ho, Jonas Schneider, llya Sutskever, Pieter Abbeel, and Wojciech Zaremba. One-shot imitation learning. In NIPS, pp. 1087-1098, 2017.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. CoRR, abs/1802.06070, 2018. URL http: //arxiv. orgo http//b//18000 06070.
Chelsea Finn, Tianhe Yu, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot visual imitation learning via meta-learning In CoRL, volume 78 of Proceedings of Machine Learning Research, pp. 357 368. PMLR, 2017.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in
actor-critic methods. In ICML, volume 80 of Proceedings of Machine Learning Research, pp. 1582-1591. PMLR, 2018
Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation. In NeurIPS, pp. 4028-40000 4028-4039, 2021.
Laurent George, Thibault Buhet, fimilie Wirbel, Gaetan Le-Gall, and Xavier Perrotton. Imitation learning for end to end vehicle longitudinal control with forward camera. CoRR, abs/1812.05841,, 2018. URL http: aariv. org/ abs 11121111111 05841.
Siddhant Haldar, Jyothish Pari, Anant Rai, and Lerrel Pinto. Teach robot fish: Versatile imitation from one minute of demonstrations, 2023.
Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. 2020.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In NIPS, pp. 4565-4573, 2016.
Jonathan Ho, Jayesh K. Gupta, and Stefano Ermon. Model-free imitation learning with policy optimization. In ICML., volume 48 of JMLR Workshop and Conference Proceedings, pp. 2760- 2769. JMLR.org, 2016.
Ahmed Hussein, Eyad Elyan, Mohamed Medhat Gaber, and Chrisina Jayne. Deep imitation learning for 3d navigation tasks. Neural Comput. Appl., 29(7):389-44)) 404, apr 2018. ISSN 0941-0643, doi: 10. 1007/s00521-007-3241-1 URL https: /tt///dooo /ooi .org 10001000//05550 000000//002100000 011-1241---11--1--
Stephen James, Michael Bloesch, and Andrew J. Davison. Task-embedded control networks for few-shot imitation learning. In CoRL, volume 87 of Proceeuiiggs of Machine Learning Research, pp. 783-795. PMLR, 2018
Jordan Juravsky, Yunrong Guo, Sanja Fidler, and Xue Bin Peng. PADL: language- physics- based character control. CoRR, abs/2301.138633 2023.
Geon-Hyeong Kim, Jongmin Lee, Youngsoo Jang, Hongseok Yang, and Kee-Eung Kim. Lobsdice: Offline learning from observation via stationary distribution correction estimation. In NeurIPS, 2022a.
Geon-Hyeong Kim, Scokin Seo, Jongmin Lee, Wonseok Jeon, HycongJoo Hwang, Hongseok Yang, and Kee-Eung Kim. Demodice: Offline imitation learning with supplementary imperfect demonstrations. ICLR. OpenReview.net, 2022b.
Kyoichiro Kobayashi, Takato Horii, Ryo Iwaki, Yukie Nagai, and Minoru Asada. Situated GAIL: multitask imitation using task-conditioned adversarial inverse reinforcement learning, CoRR, 1.b0/19 2019.
Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off- pol-cy distribution matching. In ICL.R OpenReview.net, 2020.
11

---

Published as a conference paper at ICLR 2024
Michael Laskin, Hao Liu, Xue Bin Peng, Denis Yarats, Aravind Rajeswaran, and Pieter Abbeel. Cic: Contrastive intrinsic control for unsupervised skill discovery. arXiv preprint arXiv:2202.00161, 2022.
Youngwoon Lee, Andrew Szot, Shao- Sun, and Joseph J. Lim. Generalizable imitation learning from observation via inferring goal proximity. In NeurIPS, pp. 16118-16130, 2021.
Rudolf Lioutikov, Gerhard Neumann, Guilherme Maeda, and Jan Peters. Learning movement primitive libraries through probabilistic segmentation. The International Journal of Robotics Research, 36(8) (178))877--89999 2017. doi: 10. 10.117/0073364997711111 URL https::/ddo....... org/ 10. 1171 01117/777777699777777
Evan Zheran Liu, Milad Hashemi, Kevin Swersky, Parthasarathy Ranganathan, and Junwhan Ahn. An imitation learning approach for cache replacement. In Proceedings of the 37th International Conference on Machine Learning, ICML' 20. JMLR.org 2020.
Fangehen Liu, Hao Liu, Aditya Grover, and Pieter Abbeel. Masked autoencoding for scalable and generalizable decision making. In NeurIPS, 2022.
Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. International Conference on Machine Learning, pp. 6736-6747, PMLR, 2021.
Yicheng Luo, Zhengyao Jiang. Samuel Cohen, Edward Grefenstette, and Mare Peter Deisenroth.
Optimal transport for offline imitation learning. In ICLR. OpenReview.net, 2023.
Yecheng Jason Ma, Andrew Shen, Dinesh Jayaraman, and Osbert Bastani. Versatile offline imitation from observations and examples via regularized state-occupancy matching. In ICML, volume 162 of Proceedings of Machine Learning Research, pp. 14639-1 14663. PMLR, 2022.
Zhao Mandi, Homanga Bharadhwaj, Vincent Moens, Shuran Song, Aravind Rajeswaran, and Vikash Kumar. Cacti: A framework for scalable multi-task multi-scene visual imitation learning, 12 2022.
Russell Mendonca, Oleh Rybkin, Kostas Danilidis, Danijar Hafner, and Deepak Pathak. Discovering and achieving goals via world models. NeurlPS, pp. 24379-24233 2021.
Ashvin Nair, Dian Chen, Pulkit Agrawal, Phillip Isola, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Combining sulf-supervised learning and imitation for vision-based rope manipulation, In 2017 IEEE International Conference Robotics and Automation (ICRA), pp. 1146-2153 IEEE Press, 2017. doi: 10. 101110//AAA2222777999244 201 7 7989247.
http::/ /doi. org 10. 11....11//ICRR..
Xue Peng, Erwin Coumans, Tingnan Zhang, Tsang- WeingWWeilll Lee, Jie Tan, and Sergey Levine. Learning agile robotic locomotion skills by imitating animals. 07 2020. doi: 10. 15.65557/SSS. 2020.XVI.00.
Xue Bin Peng, Pieter Abbecl, Sergey Levine, and Michiel van de Panne. Deepmimic: Example- guided deep reinforcement learning of physics-based character skills. ACM Trans. Graph., 37(4), jul 2018. ISSN 0730-0301. doi: 10. 10.1145319955173200311 URL https: /ttl:://doo.ioeo
httpp:://di..rgg// 10.
1145/3197517.3201111
3201311.
Xue Bin Peng, Yunrong Guo, Lina Halper, Sergey Levine, and Sanja Fidler. ASE: large-scale reusable adversarial skill embeddings for physically simulated characters. ACM Trans. Graph., 41 (4):94:1-94:17, 2022.
Karl Pertsch, Ruta Desai, Vikash Kumar, Franziska Meier, Joseph J. Lim, Dhruv Batra, and Akshara Rai. Cross-domain transfer via semantic skill imitation, 2022.
Bilal Piot, Matthieu Geist, and Olivier Pietquin. Boosted and reward-regularizee classification for apprenticeship learning. In AAMAS, pp. 1249-1256. IFAAMAS/ACM, 2014.
Dean Pomerleau. ALVINN: an autonomous land vehicle in a neural network. In NIPS, pp. 305-313. Morgan Kaufmann, 1988.
12

---

Published as a conference paper at ICLR 2024
Siddharth Reddy, Anca D. Dragan, and Sergey Levine. SQIL: imitation learning via reinforcement learning with sparse rewards. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30. 2020. OpenReview.net, 2020. URL https: //openreview. www.wieweeeee net
forum? nwt.oruuuuumm x xKd24twR.
Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using score-based diffusion policies. Robolics: Science and Systems, 2023.
Stéphane Ross, Geoffrey J. Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In AISTATS, volume 15 of JMLR Proceedings, pp. 627-635. Jogg. 2011.
Stefan Schaal. Learning from demonstration. In NIPS, pp. 1040-1046. MIT Press, 1996.
Lior Shani, Tom Zahavy, and Shie Mannor. Online apprenticeship learning. In AAAI, pp. 8240-8248. AAAI Press, 2022.
Lucy Xiaoyang Shi, Archit Sharma, Tony Z. Zhao, and Chelsea Finn. Waypoint-based imitation learning for robotic manipulation. CoRR, abs/2307. 14326, 2023.
Zhenyu Shou, Xuan Di, Jieping Ye, Hongtu Zhu, Hua Zhang, and Robert Hampshire. Optimal passenger-seeking policies on e-hailing platforms using markov decision process and imitation learning. Transportation Research Part C: Emerging Technologies, 111:91-113, 2020. ISSN 0968-
090X. doi: com/ /ssienee/arrrtceee///////8800000
2011000000000002001020000000 URL https: www.::/scenneiiee......
Harshit Sikchi, Qinqing Zheng, Amy Zhang, and Scott Niekum. Dual rl: Unification and new methods for reinforcement and imitation learning, 2023.
Richard S. Sutton and Andrew G. Barto. Reinforceemett Learning: An Introduction. The MIT Press,
second edition, 2018 URL html.
Umar Syed and Robert E. Schapire. game-theoretic approach to apprenticeship learning. In NIPS, pp. 1449-1456. Curran Associates, Inc., 2007.
Umar Syed, Michael H. Bowling, and Robert Schapire. Apprenticeship learning using linear programming. In ICML, volume 307 of ACM International Conference Proceeding Series, pp. 1032-1039. ACM, 2008.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin Riedmiller. Deepmind control suite. CoRR, Abs/1881.00661 2018a. URL http: htt//axiiv.oo org/ borgg/all /bs/ 1801. 00690.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin Riedmiller. Deepmind control suite. CoRR, abs/1801.00690, 2018b.
Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In Jérôme Lang (ed.), Proceedings of the Avvery-Seventtl International Joint Conference Artificial Intelligence, IJCAI 2018, July 13-19, 2018, Stockholm, Sweden, pp. 4950-4957, 4957. ijcai.org, 2018. doi: 10.24963/ 2jcai.2018/687. URL https: doi .rg 2...2.2933/000 j/aeii///jjai... .2188.01187 000018687.
Ahmed Touati and Yann Ollivier. Learning one representation to optimize all rewards. In NeurIPS, pp. 13-23, 2021.
Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? In ICL.R OpenReview.net, 2023.
Luca Viano, Angeliki Kamoutsi, Gergely Neu, Igor Krawczuk, and Volkan Cevher. Proximal point imitation learning. In NeurlPS, 2022.
13

---

Published as a conference paper at ICLR 2024
Nolan Wagener, Andrey Kolobov, Felipe Vieira Frujeri, Ricky Loynd, Ching- An Cheng, and Matthew Hausknecht. Mocapact: multi-task dataset for simulated humanoid control, 2023.
Alexander W. Winkler, Jungdam Won, and Yuting Ye. Questsim: Human motion tracking from sparse sensors with simulated avatars. In SIGGRAPHI Asia, pp. 2:1-2:8. ACM, 2022.
Philipp Wu, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Rajeswaran. Masked trajectory models for prediction, representation, and control, 2023.
Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator Discriminatorr-eightte weighted offline imitation learning from suboptimal demonstrations. In ICML, volume 162 Proceedings of Machine Learning Research, pp. 24725-24744 PMLR, 2022.
Sherry Yang, Ofir Nachum, Yilun Du, Jason Wei, Pieter Abbeel, and Dale Schuurmans. Foun- dation models for decision making: Problems, methods, and opportunities. arXiv preprint MAXi::2303.04229 2023.
Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don't t change the algorithm, change the data: Exploratory data for offline reinforcement learning. arXiv preprint arXiv:2201. 13425, 2022.
Lantao Yu, Tianhe Yu, Jiaming Song, Willie Neiswanger, and Stefano Ermon. Offline imitation learning with suboptimal demonstrations via relaxed distribution matching. AAAI, pp. 1 11016- 11024. AAAI Press, 2023.
Tianhe Yu, Chelsea Finn, Sudeep Dasari, Annie Xie, Tianhao Zhang, Pieter Abbeel, and Sergey Levine. One-shot imitation from observing humans via domain- adaative meta-learning. In Robotics: Science and Systems, 2018.
Haotian Zhang, Ye Yuan, Viktor Makoviychuk, Yunrong Guo, Sanja Fidler, Xue Bin Peng, and Kayvon Fatahalian. Learning physically simulated tennis skills from broadcast videos. ACM Trans. Graph. 42(4), jul 2023. ISSN 0730-0301. doi: 10. 145/359248. URL httpt:///doi org/ NB.1145/3592430
Tianhao Zhang, Zoe McCarthy, Owen Jow, Dennis Lee, Xi Chen, Ken Goldberg, and Pieter Abbeel. Deep imitation learning for complex manipulation tasks from virtual reality teleoperation. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 5632,, 2018. doi: 10.1 109/ICAA2201 84461299.
Mandi Zhao, Fangchen Liu, Kimin Lee, and Pieter Abbeel. Towards more generalizable one-shot visual imitation learning. In ICRA, pp. 2434-2444. IEEE, 2022.
Yang Zhou, Rui Fu, Chang Wang, and Ruibin Zhang. Modeling car-following behaviors and driving styles with generative adversarial imitation learning. Sensors (Basel, Switzerland), 20, 09 2020. doi: 10.3399//20185034
Zhuangdi Zhu, Kaixiang Lin, Bo Dai, and Jiayu Zhou. Off-policy imitation learning from observations. In NeurlPS, 2020.
Brian D. Ziebart, Andrew I.. Maas, J. Andrew Bagnell, and Anind K. Dey. Maximum entropy inverse reinforcement learning. In AAAI, pp. 1433-1438. AAAI Press, 2008.
Konrad Zolna, Alexander Novikov, Ksenia Konyushkova, Çaglar Gülçehre, Ziyu Wang, Yusuf Aytar, Misha Denil, Nando de Freitas, and Scout E. Reed. Offline learning from demonstrations and unlabeled experience. CoRR, abs/201 1. 1.18855 2020. URL htt.tt/s:/// //aassi////iaaa/aa http:////wwwww.///////////////// org//bs/200000 13885.
14

---

Published as a conference paper at ICLR 2024
Part
Appendix
Table of Contents
Additional Results and Proofs
16
A.1
Derivation of the Expression for FB with Reward pe/ (p+ po)
16
A.2 Proof of Theorem 2
16
A.3 Distribution Matching with FB: The Feature Matching Loss with R - (p
17
A.4
Derivation of the Bellman Gap Loss for Distribution Matching
18
A.5
KL Distribution Matching via FB
20
A.6 Using Universal Successor Features instead
Imitation
21
The Bellman Gap Loss and the Behavioral Cloning Loss Bound Distribution
Matching Losses
22
B Offline Baselines
24
C BFM Baselines
25
C.I. DIAYN
25
C.2 GOAL-GPT
26
C.3 MASKDP
26
C.4 GOAL-TD3
26
D Experimental Setup
27
D.I Environments
................................
27
D.2
Datasets and Expert Trajectories
................................
27
D.3
Architectures..
28
D.4
Hyperparameters.
29
E
Imitation Learning Experiments
31
Detailed View of The Results
Expert Demonstrations
31
E.2
Warm Start for FB-IL Methods
31
E.3
FB Model Quality
32
E,4
Effect of the Number of Demonstrations
33
E.5
Generalization to Different Starting Point Distributions
34
E.6
Distribution Matching: Matching the Successor Measure on Average vs al each State
35
F
Waypoint Imitation Learning Experiment
F.1
Detailed Protocol
38
F2
Algorithms and Hyperparameters
38
F33
Qualitative Evaluation
38
G Baseline Results on D4RL
40
15

---

Published as a conference paper at ICLR 2024
A ADDITIONAL RESULTS AND PROOFS
A.1 DERIVATION OF THE EXPRESSION FOR FB WITH REWARD Pe/(p pc)
In FB theory, the expression following the original training distribution p. This is not always possible or desirable. In general, if we have access rewards r(s) for states sampled from some distribution p', then the expression for becomes (Touati & Ollivier, 2021, $B.5)
(rs) assumes that sample rewards r(s) for states
Een (Cov (Bsrpp - (r(s)B(s)
(17)
which reduces to Bs-plr(s)B(s)) when p p=p'.
Taking r(s) = P. (ds) pdd and setting p' == =0p + Po) yields
-1
Po(ds)
= (Cov B)
Cov B + 2
Esmp. B(s))(s)T
E
B(s)
(18)
3~ s
((ds) + Po pd(s)
Cov B Esup, B(s))(s))
-1 / (Ep(ds) + (cs(
Pc(
(Cov B)
B(s)
2
p(ds) Pe(ds)
(19)
(Cov B) (Cov B + Esmp, B(s)B(s))) Eemp, [B(s)]
(20)
as needed.
The choice p' = 5 (p+ Pe) is equivalent to computing from a theoretical dataset that would mix the unsupervised and expert datasets. Note, though, that the final expression does not explicitly involve the unsupervised dataset, beyond the covariance matrix Cov B.
A.2 PROOF OF THEOREM 2
The loss is
LB (2) = Essrepo
||ʃ B(s) ds) - (st
(21)
= Escerpe ||ʃ
ds) 112 Exixp. (ʃ (ss) d (J B** ds)) + cst. (22)
where the constant does not depend on
Under the successor feature property (5) of FB, we have, for 8 S,
f (s, = E
80 = s, r' B(s+11)
(23)
$20
= (Cov F(s,))
(24)
where the first equality follows from the definition of the successor measure (11=*s,ds)) and the second one from (5).
Therefore, the loss Lp- rewrites as
CB- (z) Escrpe [Cov ()F =)112-22
22-22E, 2., ((Cov B)F(st,))) (J B(s) + est. (25)
Note that, for the derivation above, we do not use the full M FT Bp property (Equation 2). We only use that this property holds when integrated against B. For this, it is enough that (Cov B)F be the successor features of B, which holds by Proposition This is weaker requirement than the full successor measure equality R*= = FT Bp.
Now, in expectation over the expert trajectories (So, $1,..., St, St....),
.), we have
=E
k B(s++k+1) 1t, Te
(26)
k20
16

---

Published as a conference paper at ICLR 2024
by definition of the successor measure M So
Excrpe ((Cov (s) (S (s+,
Espepe
(CCov (CFo:))((s E
Σ
+ * B(s++x+1)
88 Te
k20
Bs Brepp E
((Cov ((Cov FF),
Σ
B($(44+11 Te (27)
k20
by properties of conditional expectations.
So the loss can be rewritten as
CB (2) Eecrop
[[ B)F(ss, 2((Cov ())
Σ
B(( + est
St, te
420
(28)
[Cor (Cov F)
Ex B((4()))+1)
+ cst St Te
(29)
for new constant that does not depend on This ends the proof.
A.3 DISTRIBUTION MATCHING WITH FB: THE FEATURE MATCHING Loss WITH R T.3
A general norm on measures is the dual L2-norm, which corresponds to taking R = L2((p) in the distribution matching criterion (10). Explicitly,
/ ((s) ||M/Plle
sup || [P r, 1,
(30)
Since f (((s) (Sg. ds) is the expected total return policy x starting at su for reward r, this norm is the worst-case optimality gap unit-norm rewards. Namely, compare the worst-case difference of two measures on all unit-norm reward functions.
The resulting loss in (12) is
(p) Earps (M (s)
(31)
This loss is tractable thanks to the following result.
Theorem 3. Assume that the FB model (2) holds. Then the quantity
Encop F(st,2) 2) (Cov (ss 2., F(ss, 2)T ?
B((++++1)
(32)
is an unbiased estimate of the loss L()(O)> (z), up to additive constant that does not depend on
So to optimize this loss, we just have to compute the discounted average of features B of states along the expert trajectory starting at each visited state st. Then we can perform stochastic gradient descent with respect to 2. The matrix Cov B is computed during pretraining.
Comparison of the L loss and the feature matching loss. The loss (32) coincides with the feature matching loss (67) when Cov B = Id. Standard FB training weakly enforces Cov B x Id by an auxiliary loss on Cov Id (Touati et al., 2023), so we do not expect a significant difference between (32) and (67).
Actually, the L loss (32) could be considered as a "canonical" version of the feature matching loss (67) that first orthonormalizes the features. The 7,2 loss is defined just by the Tp)(p) norm without
17

---

Published as a conference paper at ICLR 2024
any reference to feature basis, so it is invariant by linear transformations within the features, while the feature matching loss not.
Compared feature matching, Theorem 3 requires stronger properties of FB than Theorem 2. Indeed, Theorem 2 only requires that the FB equation F Bp = .M hold on the span of B (and would hold similarly for other successor feature models), while Theorem really requires F Bp M as measures.
Proof of Theorem 3. The proof is as follows. We have || = [(M(ds)/o(de")) p(ds') Therefore, with starting point 8, we have
||M* = (9) (
2 (M -20
+ est
2 (s. = - (s ds') (sdd (.
(s, + est (33)
where the constant term does not depend on s. If the FB model (2) holds, we have MT= Ms= F(s, 2) " B(8). Therefore,
/ (F(s,2) (()) -2 2/ F B(S) (34)
[sM* M (s)
= F(s, (Cov ()) - 2F( / (ss))ar=sss (35) (Cov
Now, in expectation over the expert trajectories (so, $1,..., st....). we have
$ B(s++x+11 Sts Te = (sss)
(36)
by definition of the successor measure M So
MM) F(s.,)) =)_ (Cov ()t, 2)), E
B(se+R+1) 8t, Te
(37)
which proves the theorem by taking expectations with respect 8t in the expert trajectories. This proves Theorem 3.
A.4 DERIVATION OF THE BELLMAN GAP LOSS FOR DISTRIBUTION MATCHING
Here we prove the claims in Section 4.3 about the Bellman gap losses.
R satisfies the Bellman equation = p >P* M So can look for the R for which this Bellman equation is best satisfied. This results in the loss (15).
We deal in turn with the seminorm I|-1l from the choice R {r= w B, E Rd ||w||| S 1} (Thm. 4), and with the full norm from R L²(p) (Thm. 5). The latter loss turns out to be similar to loss used for FB training (Touati al., 2023, $5.2), but using transitions from m. instead π.. and working with M(s, ds') instead of (M, a, ds').
Theorem 4. Assume the FB successor feature property (5) holds. Let 8t E S.
Let M and F denote the stop-grad operator over M and F (still evaluated at M M and F F).
Then the following quantities have equal gradients with respect to
The ---norm of the Bellman gap, || pa= 7 pa 7/ | at state 8t-
18

---

Published as a conference paper at ICLR 2024
The expectation over transitions 81+1 from policy Te, of the quantity
2F 2) (Cov B)
- - (((++::22) (Cov (F - -F((5)+1, (38) (Cov Theorem 5. Assume the FB model M (8, (8, de') - 2) B(8) (((s') holds. Let 81 E S.
+
Let M and F denote the stop-grad operator over M and F (still evaluated at M M and F == F).
Then the following quantities have equal gradients with respect to 2:
The (p) -norm of the Bellman gap, ||M - p 7 M (o))) at state 8t.
The expectation over transitions 8t - S141 from policy Te, of the quantity
-2F(st; 2) ((s++1)
+ (F(ss, - (+(41,, 2)) (Cov (F z) -PP(( (39) This is also the FB training loss from (Touati et al., 2023, $5.2) with transitions from Te.
The proofs are essentially the derivation used to obtain the FB loss in Touati al. (2023). the proof. we freely use the matrix and operator notation from (Blier et al., 2021) for kernels like P(s, ds') and M (s, ds'), which amounts to ordinary matrix multiplication when finite.
For Theorem 4, remember that
/ M(r. ds)w B(s)
sup weRd, 1|w1|2 |[M(+))
(40)
sup wERd, 11
w
(41)
|// M (s ds)
(42)
Therefore, the norm of the Bellman gap above is
(8(,)) - p (s6, :) PP (st, // M (s) - pmx (d -y / (ds(4112))ss) / M
MM
(s4+1,ss))(()
(43)
Now, if the FB successor feature relation (5) holds, we have
f M (st, = E B(scA
(44)
(Cov (st,
(45)
via (5). Therefore, the norm above is (Cov B)F(st, () / (dse+1 (st))++1s))B((t1111) -7 / (dse+1ll)))) (Cov B) (8)+1,
(46)
Now, usual in deep Q-learning, thanks to the stop-gre grad operator on F(S 2). we can take the expectation over 81411 P(ds++1]s [sx) outside the norm, because the difference a constant that has no gradients with respect to z. Therefore, the above has the same gradients as
E. Pre || (Cov F(F 2) B(s++1) "Coov B) (s
(47)
and rearranging yields the result.
For Theorem 5, let us start with the norm the Bellman gap of = at St.
M (st) pm= pP (-1 (-. (1:))))2
(48)
19

---

Published as a conference paper at ICLR 2024
Since || [= 2M(M// dM22 by simple computation. Therefore, the norm of the Bellman gap above is
|lM/ have |[M] M2 1
| + ||M2] Lz)0)*
|M (s6)li22))) + est 2 MM (sc, dss+1) pd pa (dst+1lst)
p(ds(+1)
-27 ʃ.
(Mr,
(1
p (ds6411))) (set1,, (s6),,
(49)
p(ds')
where the constant absorbs all terms not containing M
Now, if R** do') = =) ((s'), we have
- ( / (F(()))) 2) B(s ((s'))
(50)
(51)
| 2() B(s') ((st,2)p(ds
(52)
2) (Cov (st,
(53)
and likewise
I M A (st, dec+1) p. (dsc+1 = Pre F(si,2) 2) B(s p(ds(+1)
(54)
and
S. = (s),
pd (ds(41))))MM (sL-,. (s
p(d))
E pao (dess / F(SU12) B(5)) B(s') B(s))))
2( (BB(S) p(ds') F(s(+1,2)
(P F(s 2) (Cov F((+1,1) (55)
since B = B because B does not depend on z. Now, the Fst, (Cov B) (s which has the same gradients (((++11, 2)) (Cov 2)) -FF())) Collecting yields the result.
of the first and third term is F(s,, (
(Cov F(s
as (F(s,,
A.S KL DISTRIBUTION MATCHING VIA FB
Here we show how our approach can be extended to another divergence between M and M Inspired from the generic loss (12), consider the following loss
(xl(2)= = Exc=pa KI((M*=ss) (*)(s)
(56)
where == f + f is the generalized KL divergence between measures and g.
=
This can be estimated from expert trajectories, similarly Theorem 2. Theorem 6. Assume that the FB model (2) holds. Then the quantity
-
-Eerripe
Σn* (F'ss, =)T B(s
(57)
k20
is an unbiased estimate of the loss (x)(2), up to an additive constant that does not depend on z.
'This extends the usual KL divergence to measures which may not sum to 1, which is necessary because the model M Fr B may not be normalized. This is the Bregman divergence associated with the convex function p - plnp.
20

---

Published as a conference paper at ICLR 2024
This estimate only makes sense if the learn model n = F B only produces positive values for F
The proof is follows. Let so be any state. Then
M (so (s (s.. / M= (50 ds')
KL(M"*(su) (so)) = In M (so, ds')
(58)
(s0 ds') (s0, ds')
(su ds')
(59)
because both M and R are successor measures, so they both integrate to 1/(1 - 7). Plugging in the model M (so. ds') F(so, 2) ((s')p(ds") this equals
M (so. ds')
In F F(so: 2 B(s') pp(s))
M M** (80: ds')
(60)
In (F(F =)" B(s')) M (so, ds') + cst
(61)
where the constant does not depend
Now, given Sy and any function f. we have J J* do') = E Ei20 (( 80 Fee
This proves the claim up to replacing so with a state St ~ Pe.
A.6 USING UNIVERSAL SUCCESSOR FEATURES INSTEAD OF FB FOR IMITATION
The FB method belongs to wider class of methods based successor features measures. Many of the imitation algorithms described here still make sense with other successor feature methods, as briefly describe."
We recall the universal successor feature framework (Borsa et al., 2018). This assumes access to state features yo: S - Ed trained according to external criterion. Then Borsa al. (2018) train a family of policies (T=), together with the successor features va of yo:
(bs ao, 2) ss E (E200)(p(s441) su, ap, T2
(62)
(s) RS arg max, F(s, a,,), 2)==
At test time when faced with reward function r, USFs perform linear regression the reward on the features:
= (Eeng (()))(ss))))))) Es-p' [rs))
(63)
where p' is the data distribution at test time. Then the policy Te, is used.
The second equation in (62) is identical to FB, but the first equation is less constrained (t) = would be a solution no external criterion is applied to s, while F == B = is not a solution of (2)
as F B must represent successor measures): it is analogous to (5), which is weaker than (2).
The equation (63) for 2, is analogous to (4); beware that the variable in FB and the & variable in USFs differ by Cov B (Cov vp) factor.
Imitation learning via USFs. All of the FB-IL methods in this paper can be extended from FB other USF models, except distribution matching: as USFs don't represent the successor measure itself, they can do feature matching but not distribution matching.
We now list the corresponding losses and formulas. The derivations are omitted, since they are very similar to FB due to the analogies in the equations discussed above. For any distribution we abbreviate Cov, + = Egp (s))s((ss))))))))))
'We have focused on FB as it has demonstrated better performance for zero-shot reinforcement learning (Touati al., 2023), and is arguably better founded theoretically, by training F and B with single criterion with risk of representation collapse.
21

---

Published as a conference paper at ICLR 2024
For behavioral cloning. the loss (7) is unchanged: identify the most likely x, given the trajectory, then use TT. This applies to any method that pre-trains a family of policies (Tx)).
For reward-baste methods. plugging the reward r(-) po()/p()) (Kim al., 2022b;a; Ma et al., 2022) into (63) with p' := yields
= (Cov, s) Ep. [+] = (Cov, s0) Er
Σ
(64)
[(T)
120
Similarly. the reward r(-) = =Pol)//ol)) + pa(-)) used in (Reddy et al., 2020; Zolna et al., 2020) leads to
= (Cov, & + Covp, =p) -)=t [p].
(65)
by using p' := 2/ + Pe) in (63).
Feature matching can be done by starting with the generic loss (12) and choosing = fr w 9, E [|w||1 I 1). similarly to FB. The resulting loss
L.g. (=) == Bea-pe p. IIʃ (**) (so, ds) - ( ds)
(66)
can be estimated by similar derivations as in Theorem 2, leading to the practical estimate
(2) = Es E [((s))) Eizo 7 (((++4++1))) St, Te est.
(67)
where we have abbreviated (1t,, =) (e))) (ei) sh(s,, ae, z). Namely, this finds a that matches the successor features (s 2) to the empirical successor features computed along the expert trajectory starting at st.
Bellman gap matching can be done with the seminorm associated with R. {r w wT yp, E Rd ||l||| I 1}. but not with the full norm from R = (p) (which requires FB). This will minimize the norm of Bellman gaps for reward functions in the span of ip, using that = v if the USF model (62) holds.
Explicitly, starting with the Bellman gap loss (15), ||M - p M** R. with transitions from the expert trajectories, and plugging this choice of R together with the USF property = = vb, leads to the loss
C Bel)((())) = Eserppe (xb 2)J (((+tt))))))
+ (1b(1),2) - 2)) (N(s), 2) - 778(5111, 2]]] + est (68)
analogous to (16).
Finally, waypoint imitation can be done in USFs by selecting goal state sttk, and putting single reward at $1+* in the USF formula (63), which yields
Zt (Cov, 1p) (p((++)))
(69)
analogously to Sec. 4.4, then using T. at time t.
A.7 THE BELLMAN GAP LOss AND THE BEHAVIORAL CLONING LOSS BOUND DISTRIBUTION MATCHING LOSSES
Any method that provides a policy close to Te will provide state distributions close to d=, result, so we expect a relationship between the losses from different approaches. Indeed, the Bellman gap loss (15) bounds the distribution matching loss (12), and the BC loss bounds the KL version of (12). This is formalized in Theorems 7 and 8, respectively. Theorem is analogous to the bound between Bellman gaps and errors on the Q-function.
Theorem 7. Let pe be a stationary distribution of the expert policy тe- Then the Bellman gap loss (15) on successor measures bounds both the feature distribution matching loss (12) and the original feature or distribution matching loss (11). Namely, for any choice of R.
Eaoop. =² (so) - Esoop. Pe M (s0))2
E Esg Pp ||M - M (50,
1
Eomp. 11 p P T (s)] (70)
(1 7)2
22

---

Published as a conference paper at ICLR 2024
Theorem 8. Let Pe be a stationary distribution over states and state-actions of the expert policy ne. Then
1 - 1-7 Benp, (-1s
Bawpe
(71)
y Blaa) (als) + est
(72)
where the constant does not depend on z.
Proof of Theorem 7. The proof is as follows. The first inequality is by convexity of the norm. For the second one, we have
(00) M (50) == (s0: -) - (Id -7 -) 1 pd (80.
(73)
= (1a-) pF. )M (
:) - - (80, )))
(74)
and therefore
Esgrre 1s0=(s00)) - (500 )|
(75)
1.0rd- (Tdd- -ppp--p-o--- (s0, ds) ((Id P - pT (s. :))
2
= Esooree
(76)
R
and since (Id - is the successor measure of TT. it is measure with total mass 1/ (1
so the integral under (Id -y px=))= can be rewritten as an expectation under (1 - -Id -7 P'e
1 (1 7)2 Beorupe (1- -Ppe))d--Pp---------- (so. [(Id -PP (s,-) pm (8, )11
(77)
(1 71- Bsorp
(s. - p=* pT M= (s, )l
(78)
(by convexity)
E. Bempe (s, ()) -7p" (s, -l
(79)
(1 72 where the last equality uses that Pe is a stationary distribution of ne, which implies that the marginal distribution of s in the above is (1 7) S.g (Id- (so. ds) (1
7) Sag Po(dso) Ei 3'(T ds) = (1 - 7) Ei nt So po(dso) ( ds) = (1
7) Ei (dss) (ds) since Pe is invariant under p.e.
Proof of Theorem (the BC loss is an upper bound of the KL successor measure matching loss). The proof goes as follows. Denote by ~=00² the probability distribution of the state s when starting at 8n and following policy æ. Denote by =0,0tt the probability distribution of the trajectory (8o, 40; 81, oi,...,, 4, 8t) when starting at 80 and following policy æ.
Denote by E, the expectation under a random variable with geometric distribution of parameter 1 y starting at 1. By definition,
(1 (1-))M² = EE, 00.6
(80)
Then.
KTL((1 - -7) M² KT ( n || E, av.
(81)
< E, KL s0 || 700.0)
(82)
< E, KL m0t *o.0t))
(83)
Te (S0, ao. .8)
88)
in E, E
| 10, Te
(84)
T(50, Co .......... s1)
E, E
In (ax(sk) Inx((al33)) 30. Te
(85)
23

---

Published as a conference paper at ICLR 2024
Now, when integrated over so Pe with Pe stationary distribution of Te. then each state sy is itself distributed according to Te. Thus, the above is equal to
Esorpe KL((1 - () (*)))((()))))))))))))))))))))( lI (1 7) (so)) - (B, Eawp, (-( (+1s))
(86)
where (E, t) accounts for there being t terms in the sum. Under the law of t, we have E t 1-?"
Finally, Hspp, (L(Te(-Is) TT. (-[s) is the behavior cloning loss up to the entropy of Te, which does not depend on 2.
B OFFLINE IL BASELINES
In this section, we provide more information about the baselines we considered in the paper. Refer to App. E for a complete view of the results. App. G we confirm that our implementation of the baselines matches results reported in the literature.
Soft Q Imitation Learning (Reddy et al., 2020, SQII.) a simple imitation algorithm that can be implemented with little modification to any standard RL algorithm. The idea is to provide constant reward r = to expert transitions, and a constant reward r = 0 to any transition generated by the RL algorithm.
SQIL has been introduced and shown to perform well in the online setting. In our experiments we use straightforward offline adaptation of the algorithm: we provide = 1 to expert transitions and r = 0 to the offline unsupervised transitions and use SAC as RL algorithm. We use balanced sampling defined in the original paper, i.e., the batch comprises equal number of expert and unsupervised transitions.
Since SQIL underperformed in several tasks, introduced TD3 Imitation Learning (TD3IL). TD3IL uses {0, 1} reward for unsupervised and expert samples as in SQIL, but uses TD3 as the offline RL algorithm, and does not use balanced sampling. We tested different ways to construct the batch provided to the agent during training. These methods included sampling expert transitions based on their proportion relative to the unsupervised transitions and sampling a fixed ratio of expert transitions. We found the fixed ratio strategy to be consistent across different number of expert trajectories, thus we used this approach in the experiments.
We also tested TD3 with soft variant of the (0. 1} reward used by SQIL and TD3IL. In particular, Luo et al. (2023) suggest to use optimal transport to compute a distance between expert and non- expert transitions that is used as a reward function. This reward can subsequently used with any RL algorithm. For the latter, we used TD3 our experiments since proved to be the most consistent. We called this algorithm OTRTD3-
Offline Reinforced Imitation Learning (Zolna et al., 2020, ORIL) trains a discriminator D(-) to separate samples from expert and unsupervised trajectories. Then, it trains offline RL agent by annotating transitions using the learned reward function r(-) = log( (--) + 1). We pretrained the discriminator using gradient penalty but we did not use positive-unlabeled learning since it did not improve performance in our tests. We trained three variants of the discriminator, using observation, (observation,action and (observation, next-observation).
Similarly to ORII,, *DIC algorithms (Kim et al., 2022b; Ma al., 2022; Kim et al., 2022a) use discriminator reconstruct the reward function. The main difference that they aim to reconstruct a reward function of the form r(-) p.+p
_a__. while ORIL targets a reward r(.) = A We use the same discriminator training procedure but we relabel the transitions using the parametric reward function r(-) = log (D(+) log(1 D())). These methods leverage a regularized RL approach for training the policy. IQ-Learn (Garg et al., 2021) is non-adversarial II method based on the MaxEnt formulation (Ziebart et al., 2008). Similarly to ValueDice (Kostrikov et al., 2020), the idea of IQ-Leamn is to transform the problem over rewards to problem over Q-functions. Opposite to *DICE algorithms, IQ-Learn does not require to explicitly train a discriminator to recover a reward function.
Discriminator Weighted Behavioral Cloning (Xu et al., 2022, DWBC) proposes to approach offline IL through weighted behavior cloning, where the weights provided by a discriminator. This is way of leveraging unsupervised samples in BC, on top of expert samples. As suggested in the paper, use discriminator conditioned on the policy learned through weighted BC.
24

---

Published as a conference paper at ICLR 2024
C
BFM BASELINES
In this section we provide a description of the behavior foundational models used in our experiments.
C.1 DIAYN
DIAYN (Eysenbach et al., 2018) is a skill discovery algorithm, commonly used in unsupervised RL. It learns set of parametrized policies (T=) by maximising the mutual information I [s: =) between the state produced by policy TN and the latent skill drawn from prior distribution of skills = ~ p(z). To obtain a tractable approximation of [(s; =), Eysenbach et al. (2018) introduce an approximate skill posterior (skill encoder) g(> 8) and use the following variational lower bound:
[(s; H()) [ |s)
(87)
H()) + -pp(s) Bendts [l 8)]
(88)
2 H()) + Ep(xs) Bend" [log ((= s)]
(89)
where Hz) and (( ()) are respectively the entropy of the prior distribution p(z) and the entropy of the conditional skill distribution p(= | 8). The latter approximated by the skill encoder q(> s). Similarly to Hansen et al. (2020), we model the latent space as d-dimensional hypersphere {=: || 1} and p(z) uniform distribution on the hypersphere. Since latent variables live on the hypersphere, we model the skill encoder as von Mises-Fisher distribution with a scale parameter
of 1:
q(= 8) x exp(= T (sx))
(90)
where yp: S - Rd is feature map restricted to the unit-hypersphere, i.e., ||<(())] = 1, 1, E S.
In practice, we train online both the feature map yo and policy nz: the feature map o learned by maximizing the log-likelihood objective: max Ep(s) [=- (=)], while n₃ is trained to maximize the intrinsic reward FDMAYN 2) T (ss) by using z-conditioned TD3.
Our first attempt to vanilla online train DIAYN was unsuccessful and leads to near-zero performance. This is consistent with findings of prior work (Laskin et al., 2022) that DIAYN is not able to learn diverse skill on DMC environments due to absence resetting when the agent falls. Therefore, we decided incorporate two components in our training that boost performance:
Posterior Hindsight Experience Replay (Choi et al., 2021, P-Her): instead of considering the actual that generates the state it consists in relabelling = for some states 8; by sample drawn from the skill encoder ( ${+k), where st+* is a future state in the trajectory of In practice, we set 2 = ((*++*), which is equal to the mean of distribution.
Exploration bonus: To incentivize the agent to learn diverse skills, we add a k-nearest neighbors- based entropy reward similarly to (Liu Abbeel, 2021): Texplora (8)
in [1 + =..EnN((((a)) Ilı(s) xill) where 2, sp(s;) for a mini-batch of states {s:). The policy is then trained to maximise the reward r(s, 2) (DIAYN(S, Xexxploro (s), where is an exploration coefficient.
DIAYN for Imitation Learning. We devise here three different II. methods based on a pre-trained DIAYN model:
Behavioral cloning (BCdiayn): we look for the policy ny that best fits the expert trajectories in term of its likelihood, by minimizing the loss:
min E, Σ In (a,
(91)
maximization (ERDiAYN): we infer the latent variable 2 by maximizing the mutual information I(x; In practice, we maximize instead the (tractable) variational lower bound.
max Bx-pa Ba-p. |s] max Eanpa (2 (2Tx(s)) max Eawpa [s(ss]
(92)
25

---

Published as a conference paper at ICLR 2024
The above maximization problem admits a closed-form solution: = Bgypp [p{ss] which consists simply in averaging the features of the expert states. This final formula is similar to (8) for ER hence notation ERDIAYX- Note that this approach was proposed by Eysenbach et al. (2018) for descrete skills (see their Appendix G). Here we report the variant for continuous spaces. Goal-based Imitation (GOAL-DIAYN): Consider one single trajectory т. At each time step I, can use the DIAYN pretrained behaviors to reach a state St44 slightly ahead of 8t in the expert trajectory. Specifically, set 2; <((t+*), which corresponds the mean of the skill encoder distribution for St+,, and use the action given by Tz, (st).
C.2 GOAL-GPT
GOAL-GPT (Liu al., 2022) goal-conditioned transformer- peasoreee-bbseed auto-regressive model x, trained offline using a behavior cloning objective. At train time, given an offline dataset of trajectories, we sample sub-trajectories {st, aty..., .....,,, $14R) consisting contiguous state-action pairs, we relabel the last state goal g S1+k: and then minimize the following objective:
E, [In (at..... detk (8t:9) ..., (81+): E,
In m (ar+z (st,9) .... (8L+k:
(93)
The last equality holds since the model uses causal attention masking.
We can use GOAL-GPT to perform goal- goal-based imitation. Given one single expert trajectory T ......., s..), we divide the trajectory into segments of equal length For the segment (sf..... s.t)). we set the goal to the last expert's state, = +f+ and we execute the A: actions predicted by for all i E (1..... k}. Here (st...., st+i) is the history of the last states generated while interacting with the environment to imitate the expert's trajectory.
the model: (a;,..., at++ ()) ..., (st+i, g))
C.3 MASKDP
Masked Decision Prediction (Liu et al., 2022, MASKDP) is a self-supervised pretraining method. Unlike autoregressive action prediction used in GOAL-GPT, it employs masked autoencoder to state-action trajectories, which randomly masks states and actions and is trained to reconstruct the missing data. It uses encoder-decoded architecture h. Both encoder and decoder are bidirectional transformers. The model is trained to reconstruct a sub-trajectory given masked view of itself, i.e., i =immasked()) T.
At train time, given a sub-trajectory 7 of state-action pairs, we apply random masking on states and actions independently. The encoder processes only the unmasked states and actions. The decoder operates on the whole sub-trajectory of both visible and masked encoded states and actions, while replacing each masked element by a shared learned vector (called a mask token). The overall model is learned end-to-end using reconstructing loss (mean square error).
We can use MASKDP perform goal-based imitation as follows: Given one single expert trajectory (sf...., ), ), at each time step f, we use the MASKDP model to predict the actions necessary to reach the goal + To this end, we fed the model by the masked sequence (st. ................................ s++ of length k where denotes the missing element, then, we execute only the first action predicted by the model.
C.4 GOAL--TD3
For Goal-TD3, pre-train offline goal-conditioned policies T[a s. using the sparse reward r(s, g) = Tf]ls g|1 I #} for some small value of e. Training uses Hindsight Experience Re- play (Andrychowicz et al., 2017, HER).
At test time, Goal-TD3 can implement goal-based IL: at each time step t, we select the policy (a 8t, x(+)) where the goal 844 is the state k steps ahead of 8, in the expert trajectory.
26

---

Published as a conference paper at ICLR 2024
Figure 5: Maze, Walker, Cheetah and Quadruped environments used in experiments.
D EXPERIMENTAL SETUP
In this section we provide additional information about our experiments.
D.1 ENVIRONMENTS
All the environments considered in this paper based on the DeepMind Control Suite (Tassa et al., 201 8b). In total, we have environments and 21 tasks.
maatt-mass-maze: a 2-dimensional continuous maze with four rooms. The states are 4- dimensional vectors consisting of positions and velocities of the point mass (x, y, Uz Ty). and the actions are 2-dimensional vectors. The initial state location is sampled uniformly from the top left room. We consider 8 tasks: four goal-reaching tasks( reach_top_left, and reach_bottom__e reaching goal in the middle of each room described by their (æ, y) coordinates, 2 looping tasks encourage the agent to navigate trough the all rooms while drawing two different shapes (square shape close to maze's borders for the squar re task, diamand shape for the d amand task), fast_slelcec ow task encourages the agent to loop around the maze while maintaining large velocity when moving horizontally and small velocity when moving vertically, and finally, consists in reaching the bottom left room following the long path (by penalizing the agent for moving counterclockwise).
Point-
seach_top_rigee
seach_bottom__e
consist in
Walker: a planar walker. States are 24-dimensional vectors consisting positions and ve- 6-dimensional vectors. We consider 5 different tasks: wa walker_stand reward is a combination of terms encouraging an upright torso and some minimal torso height, while wa ker_wall and wa 1 ker_run rewards include component encouraging some minimal forward velocity. wa kerkeer_f flip reward is a combination of terms encouraging and upright torso and some mininal angular momentum while walker flip reward encourages only to have some minimal angular momentum without constraining the torso. Cheetah: a running planar biped. States are 17-dimensional vectors consisting of positions and velocities of robot joints, and actions are 6-dimensional vectors. We consider 4 dif- ferent tasks: cheetah_we k and cheetah_st
locities of robot joints, and actions
rewards are linearly proportional to forward velecity up some desired values: 2 m/s for walk and 10 m/s for run. Simi- larly, walker_weer _alk backward. and walker run backward rewards encourage reach- ing some minimal backward velocities.
Quadruped: four-leg ant navigating in 3D space. States and actions are 78-dimensional and 12- dimensional vectors, respectively. We consider 4 tasks: quadruped ana reward encourages an upright torso. quadruped raaedwwaewwala _aak and quadruped_ sruued_ruuuu run include a term encouraging some minimal velecities. quadruped_wad walk includes of the center of mass. quadruped_ encourages some minimal height.
encouraging some minimal height
D.2 DATASETS AND EXPERT TRAJECTORIES
We used standard unsupervised datasets for the four domains, generated by Random Network Distil- lation (RND). They can be downloaded following the instructions in the github repository of Yarats et al. (2022) (attps:/ ///gttttuu////// com/ denisysaa denisyyratts exoorl). We use 5000 trajectories for locomotion tasks (walker, cheetah and quadruped) and 10000 trajectories for point-mass-maze
27

---

Published as a conference paper at ICLR 2024
Expert trajectories for the 21 tasks are generated by task-specific trained TD3 agents. We train TD3 online for locomotion tasks. We train TD3 offline for the navigation tasks since the maze's ss unsupervised dataset has good coverage to learn the optimal behaviors.
D.3 ARCHITECTURES
Forward-Backwand
The backward representation network B(s) is represented by a feedforward neural network with two hidden layers, each with 256 units, that takes as input state and outputs a d-dimensional embedding. The output of B can be either L2 -2-projected into the sphere radius vd batch- normalized. Using Batchnore makes FB invariant to reward translation since ensures [ss]] 1.
For the forward network F(s, 0, 2), we first preprocess separately (s, a) and 2) by two feedforward networks with single hidden layer (with 1024 units) 512-dimentional space. Then we concatenate their two outputs and pass it into two heads of feedforward networks (each with one hidden layer of 1024 units) to output a d-dimensional vector.
For the policy network (s, 2). we first preprocess separately and 2) by two feedforward networks with one single hidden layer (with 1024 units) to 512-dimentional space. Then we concatenate their two outputs and pass it into another one single hidden layer feedforward network (with 1024 units) to output to output d -immennsioa vector, then we apply a Tanh activation as the action space is 1),
For all the architectures, we apply layer normalization (Ba et al., 2016) and Tanh activation in the first layer in order to standardize the states and actions. We use Relu for the rest of layers. We also pre-normalize <- v T in the input of F, and v.
For maze environments, added additional hidden layer after the preprocessing (for both policy and forward) as it helped to improve the results.
Imitation Learning Baselines: For all IL baselines, we use feedforward neural networks with two hidden layers, each with 1024 units to represent actors, critics and discrimintors networks. We apply a layer normalization and Tanh activation in the first layer in order to standardize the states and actions.
BFMs: For DIAYN, we have three networks:
The policy network I (s, 2) is similar that of FB: it preprocesses separately and (s,2) by two feedforward networks with one single hidden layer (with 1024 units) to 512-dimentional space. Then we concatenate their two outputs and pass into another two-hidden-layee feedforward network (with 1024 units) to output to output a dagddimensional vector, then we apply a Tanh activation as the action space is [-1, 1]da.
The critic network ((s,, a, 2) is similar to F in FB but with scalar output: it preprocesses separately (s, a) and (s.2) by two feedforward networks with one single hidden layer (with 1024 units) to 1-5512222 512-dimentional space. Then it concatenates their two outputs and pass it into two heads of feedforward networks (each with two hidden layers of 1024 units) to output a scalar.
The skill encoder network sp(s) is represented by feedforward neural network with two hidden layers, each with 256 units, that takes as input a state and outputs a d-dimensional embedding. The output of yp is L2 12-projected into the mnithhyprrspher.
For COAL-TD3, we use the same architectures for both policy and critic networks as DIAYN, and replace the latent skill by goal state.
GOAL-GPT uses auto-regressive transformer architecture via causal attention masking with attention layers. Each attention layer has 4 attention heads with 256 hidden dimensions.
MASKDP uses uncoder-decoded architecture. Both encoder and decoder are bidirectional transformers with respectively and 2 attention layers. Each attention layer has 4 attention heads with 256 hidden dimensions. More details about the architectures for both GOAL-GPT and MASKDP can be found in (Liu et al., 2022)
28

---

Published as a conference paper at ICLR 2024
D.4 HYPERPARAMETERS
Table Hyperparameters used for FB pretraining.
Hyperparameter
Walker
Cheetah
Quadruped
Maze
Kepresentation dimension
50
100
Batch size
2048
2048
1024
1024
Discount factor Yv
0.98
0.98
0.98
0.99
Optimizer
Adam
Adam
Adam
learning rale of learning rale of B learning rale it Normalization
10
10-
10-
10-4
10
10
10-4
10 10-6
10
10
10-4
L2
Batchnorm
Momentum for larget networks
0.99
0.99
0.99
0.99
Siddev policy smoothing
0.2
0.2
Truncation level for policy smoothing
0.3
0.3
0.3
0.3
Regularization weight for octhonomality
Table 2: Hyperparameters used for IL baselines.
Hyperparameter
Walker
Cheetah
Quadruped
Maze
Common
Batch size
512
512
512
512
Discount factor ?
0.98
0.98
0.98
0.99
Optimizer
Adam
Adam
Adam
Adam
learning rate
10-4
10
10
10 4
Momentum target networks
0.99
0.99
Stddev for policy smoothing
0.2
0.2
0.2
0.2
Truncation level for policy smoothing
0.3
0.3
D.3
0.3
Discriminator' training steps
5 10
5 10
5 10
10
ORT.
Gradient penalty
Positive-unlabelod coefficient
0
SMODICE
Gradient penalty
Divergence
x²
2
x
2
LOBSDICE
Gradient penalty
10
10
2
10
Divergence
DEMODICE
Gradient penalty
2
Divergence
KL
KL
KL
KL
TD3IL
Sampling
0.002
Table 3: Hyperparameters used for DIAYN and GOAL-TD3.
Hyperpurameter
Walker
Cheetah
Quadruped
Maze
Common
Batch size
512
512
512
512
Discount factoor y
0.98
0.98
Optimizer
Adam
Adam
Adam
Adam
Momentum for target networks
0.99
0.99
0.99
0.99
Stddey for policy smoothing
0.2
0.2
0.2
Truncation level for policy smoothing
0.3
0.3
0.3
0.3
DIAYN
actor learning rale critic learning rate skill encoder learning
10
10
10-5
10
10
exploration coeflicient
0.1
1
ratio
Larent dimension
50
25
25
25
GOAL-TD3
actor learning rate
10
10
10
10
critic learning rate ratio
10
10
10
10
Table 4: Hyperparameters used for GOAL-GPT and MASKDP.
Hyperparameter
Walker
Cheetah
Quadruped
Maze
Common
Batch size
512
512
512
512
Optimizer
Adam
Adam
Adam
Adam
learning
10
10-4
10
10-
MASKDP
conteat size
64
64
64
64
GOAL-GPT
COnteNt size
32
29

---

Published as a conference paper at ICLR 2024
Hyperparameter finetuning. We finetune the hyperparameters of FB models for each domain by performing hyperparameter sweeps over the batch size, learning rates, representation dimensions, orthonormality regularization and normalization of backward output. We evaluate each model using its reward-based performance on downstream tasks for each domain. We select the hyperparameters that lead to the best averaged performance over each domain's 's tasks and over three random seeds. We re-train the final FB models with the selected hyperparameters for 10 different random seeds.
For imitation learning baselines, finetune each baseline's hyperparameters for each domain by performing hyperparameter sweep on a representative task (e.g.. wa walker_wald for walker and cheetah_wa walk for cheetah). We did not do hyperparameter sweeps for each baseline and task, since this would have been too intensive given the number of setups.
For the other BFMs (DIAYN, GOAL-TD3, GOAL-GPT and MASKDP), we finetune each of them by performing hyperparameter sweep on their goal-based imitation performance in representative task for each domain. We select the hyperparameters that lead the best averaged performance over three random seeds. We re-train the final BFMs models with the selected hyperparameters for 10 different random seeds.
This results in one hyperparameter tuning per domain (not per task), both for BFMs and for the IL baselines. Importantly, note that the BFM tuning shared by all algorithms using given BFM model (e.g., ERFB. BBELLFB. ... for FB), while the IL baselines have separate tuning per algorithm.
30

---

Published as a conference paper at ICLR 2024
E IMITATION LEARNING EXPERIMENTS
In this section we report the complete set of results for the standard IL protocol described in Sect. 5. We also present the additional experiments we conducted to assess the performance of FB-IL methods and baselines.
E. DETAILED VIEW OF THE RESULTS WITH K = EXPERT DEMONSTRATIONS
We start providing results for each task. Tab. 5, and contain the cumulative reward obtained by the IL policy recovered by each algorithm. The protocol is the same as in Sect. 5. We further report the performance of the expert agent, expert performance is computed over the 250 trajectories we generated for the IL. experiment.
In Tab. 8 we further report the average time required by the methods to compute the IL policy.
Maze
Algorithm
fast slow
loop
reach bottom
left
reach bottom left long
reach bottom
reach top left
reach top right
square
451.7(00.a
506.9g 7)
166.9ras
263.4:a..44
167.5cs o)
476.3m 9)
452. 22..6is
757.8(sw.)
DWBC
327.9(42.2)
(31.))
273.9
120,5c20.
453.3
205.1
587,0
ORIL
4.6::4::
42. (2.5)
0,0ic.o:
191.7
0.0raun
632.6Ga7
0.0ja.m
143.5/2..)
ORIL-O
144.8
524.1
156.1
244.408
298.1
952.2:5.0p
785.61184
390.8
ORII-0O
148.96-
488.9
715.3 9)
218. 1920.))
237.0(sna)
950.4:10
740.7 740.7
417.2 417.2
OTRTDa
121.7(w0.0
676.4
486.8
361. 8/s.00
274.0
401.5or
549. 80088oss
640.0
SQIL
700.0ts4.
7(1.0:s4.0)
(9672)
533. 9:a1.11
5-18,5can Sc4n.7
926.0
739.9 9:
69-1.5y00.8)
TD3IL
104.9:28.0
752.2
830.5,0.2
-11.1(5.7)
608.8
963. 4:1
829. (.0.6)
244.3:63.6)
DEMODICE
107.4ç10.
20.9(2.5)
-928.3366.5)
19.9
111.2:449
26.2
416.4(13-)
SMODICE
117.16.0
442.7
105. 2:0.
1340. an
88. 0:
85.9
115.5
296.7 7(x
LOBSDICE
86.70.4
109.6
15.0
2018. 20018.2/ss.s)
12.2 12.2
96.0c1.3
16.38 16.5to.s)
293.9:a.))
IQLEARN
48.0n5.3)
148.2(26.5)
0.0io.o:
1966,8rr.7
0.1caa
42. (2.1
0.2ba.a
114.4
BCDJAYn
16.0
142.3
545.8 845.8121.3)
224. Gica.a
211.4
737 737.5
584.6 1
116.1:4.5)
ERDAY
0:4
87.2:7
517.7(18
(38..8
568.5
550.9
46.2(2.m)
GOALDLAV'
215.2
388.6
580.8(
47.8
299.4cm
585.5ms
619.1
414.3
GOAL-TDS
821.3
800.9
781. 5
173. 121.5
335.4
918. 6
700.1:s.e)
90116::14)
MASKDP
692.4
763.8
592.0
461. 5:12.0
293.8:22.)
913.2:s.a
654. (/p.3)
821 221.3
GOAL-GPT
710.8
843.2
713.3:2 32..9
606.8tam7
171.6rs4.a)
665. D
716. 8:2.8)
840.8(177
392.8m..
846.8:549
734.5
1560.4(1..))
548.5(12.0)
962 962.4(e-a)
758. ((s.0)
8T3. (()0)
ERyB
34.
697. 60
817-4(o.4)
-7.33 1)
527.7(18)
956.5k1 ()
819.5
771.9(s.)
RERFB
281.8
637, 8(13.))
776.
317.7(mu.)
570.6
915. 915.5/220
760. 160.1
718. A11..)
BBELLPR
236.2
648.8115.0)
491.8:1977
53.99a1.7)
313.3
776.5/15m7
112.4n1..
445.1 (15.1
FMFB
221.2
598.6 611.6
356.5
-129.9tm
178.8(ss.))
754100mss4
432
413.9
FBLOsSpa
219.418.4)
638. O1140)
377 0110.;
-20.7:33.4)
375.5
821.911..)
297 (:1.3)
473 (73.5:14.3)
DMFB
229.8
169.1a
277.7
746.9119.)
271.5
447 447.81148)
GOAL
572.9g
799.7(1
543.9(9m
617.6cc.
954.6(1.0 o)
703.7(17)
885-5
Expert
935.31e.e)
912.410.5)
813.4:1.3)
080. 4cs.
710.311.0
949.3:1
813.0(1
952.9mas)
Table 5: Cumulative reward for each task the maze environment, averaged over repetitions. Experi- ments are done with K = 1 expert demonstrations. Standard deviation is reported in parenthesis.
E.2 WARM START FOR FB-IL METHODS
As mentioned in the text, some FB methods require a gradient descent over the policy parameter and we initialized the gradient descent with "warm start", setting the initial guess 2g the value (8) used in ERFB. which can be readily computed.
Fig. illustrates the performance of BCFB and BBELLFB with and without ERpB warm-start averaged over all environments and tasks. While BBELLFB is relatively robust, BCrB performs poorly when the policy embedding z optimized starting from a random point on the unit sphere. While BCFB benefits from warm-start, notice that the initial value 2p obtained from ERpB is not stationary and the BC loss keeps decreasing over iterations and eventually converges to a different policy æ. This can lead to policies with significantly different behavior as illustrated in Fig, 6. While BCFB without warm start tries to imitate the expert and fails in reproducing the trajectory, the policy returned by ER reaches the goal but takes different path w.r.t. the expert. On the other hand, BCFB with warm start successfully shifts the initial ERpB policy to better replicate the expert actions and eventually reproduce its trajectory.
31

---

Published as a conference paper at ICLR 2024
Walker
Algorichm
run
spin
stand
walk
BC
28.3
26.3:1.0
21.5a..3
166.7
27 2r.2
DWBC
67.2
27.6:1
190.3
191.0
39.0gz.9
ORIL
a48.0casa)
320, 2:203:
900.1 (s6.4)
886.9
567 3c
ORIL-O
743.0
262.4.22.4
9.0
34..1
114.7(:0..))
ORIL-OO
92.2
122.5
9.3
335.7 7
72.8ip.5)
OTRTDa
27.4:1
183.5
285.8
64.5
SQIL
561.0
309.10a..
967.5
807. (:500)
681 681.4.66.9
673.7
293.4
973.0
876.5 9
736.0
DEMODICE
245.111.3)
87 (
17-4.2ts.1
389.7 389.7wa))
(
SMODICE
241.3
89.5
387 D/o.5)
196.2(c.5)
LOBsDICE
243.6;0.4)
89.1:0-3)
481,9(a.))
387 (3.33
19.4
IQLEARN
34
23.1:004
57.6
139.8
B BCDAYN
36.6
23. 23.5ç0.a)
17.5
195.6;444
33.3
ERDAAN
77.0ka.)
25.0:1.0
89.8
151.7
104.8
GOAT GOALDEAYN
104.6:a.2
26.7:0.7
45 45
145.9
131.56
GoAL-TD3
593.2
216.8
859.8(a.a)
908. Giz.s)
3
MASKDP
65. T(2.4)
67. (:.6:
116.0r4.)
162.5/1.))
32,2/0.3)
GOAL-GPT
253.4...)
102.9
375. 1ea.4)
406. .
214.8
BCF3
(79.7
262 222.5)
793. 6
742.9kva (
891 001.90
ERen
352.0)a.)
222.3)
814.u/77
72..T(s.))
836.3
RERyB
553.1
343.915.3)
000.3
672.5
735.4(10.5)
BBBLLyaa
612.6
322.1(4.9)
896.0012.3)
T(0u)(
T89 21.90
FMyH
606.5:a.a)
228.1
836 836.4411a.5)
706.3cea)
808.917.6)
FBL.0SSFB
64
282.7:5.9)
910.6(11.3
(1a.)
(11.1
DMFB
61-.1
256.0
824. 324.6(1m77)
60..3
832.8
GOALLB
593.6
275.4
903.7
715.4
820.3
Expert
977.2.0.))
545. (0.7)
(80.2
388.5504))
(78.7
Table 6: Cumulative reward for each task in the walker environment, averaged over repetitions. Ex- periments are done with K = 1 expert demonstrations. Standard deviation is reported in parenthesis.
Cheetah
Quadruped
Algorithm
run
run backward
walk
walk hackward
jump
nan
stand
walk
BC
62.2:14)
88.66c0.8;
237 U1a.44
675.3
156.22cu
197.
429.8:42.
295. 3<39
DWBC
63.8
S8.6
247.9
6
161.1cc4
195.6
444.3
292. Oc2000
ORIL
200.0m..)
365.6is.s)
655.2
952.77 p..))
801.022a
571.4
663.6
763.
ORIL-O
9.6cu.6:
31.0rzo
421 1
823. 1)
680.7:5.3
454.1
968.4
555, Scaa
ORIL-OO
103.5
243.5 243.5:30.2)
459.9
820.1
677.5
958 958.2:14
OTR
91. 91.2:a.a)
107.5 (.1..)
515.4
681.6 (18.))
363. 363.2
31.6:a.0
277.1:a1.))
61.6 (1..66..)
SQIL.
194.7
187.6 8:00
120.4
560. (005.
266. 1099
473.5
361 363.3
273.2
TD3IL
214. 2inaa
347.9
861.3
968.8:a.a)
808. 2
552..6
(56.1
813.8
DEMODICE
39.6:0.)
54.4()
197.3:10
250.9
154,1ga
(1
1
114.1
SMODICE
43.1
56.9to.s
205.5/2.1
256.7
139.0
97. 0
205.9
109.4.2.0)
LORSDICE
47.2
53.6
234.9c
225.2/sa)
1 5)
m2
242.7(a.0
123.60-0
IQLEARN
46.5
(5,2
320.9
772.5
9B.8m
71.3
242.730.5
237.0k6.0)
0.8
4.900.0
4.70.20
153 4.
169.0
272.3 (:s.:
147700t.7
EReHAVX
1.6ca.
1.6
8,2.0.2)
7.610.3)
148.0
243.9ka.11
124.0h2.)
GOALDIANN
17.2
51.3psa
376.5 57)
199.5 199.5
117.40.5
125. 2/244
196 196.0
157.7c4.
GoAL-TD3
83.5ga..)
171.3
6535.1/7.0)
732,4 418.5
(26.6)
760.7
MASKDP
63.5
19.1n. 19..0..)
318.7:7. 7:
831. (18.))
150.1
90.2;22::
299 0
161.0ca.5
GOAL-GPT
49.2)o.m
97.411.3)
811.8
675.22
543.0
638.7
320.9h..)
BCra
176.6
184.5
6446.0gn. 19
629. (24.4.g4.1
784.2 (
332 Aro
967.3cx.1
800.2
EReB
303.4:s.)
207. Aia.7)
888. T(s.))
885.91 885.9112.0)
798.5;444)
137 3co.s
170.4
703.4mm7
RERFB
252.0
228, 2(4.0)
845.8 So
854.8 $54.8
632. (32.2(1w.))
393.8 m
864.4/18 (:444)
649. 000.60000
BBELLys
305.0(2.7)
231.8is.8
674.99us.
629.0112.)
629 8(12
338.3
9:11.3ra.4
762. 9:2.:
FMFB
290.9
210.3
(62.1
834.4
534. ((449)
241.0ct.0
782.3:10.4)
612. 112211113
FBLOSSyn
299.3:s.)
237. 7
632, 1ns.0
821.6r1s Gr1s.)
613 ((1..8
323. 323.0pa..
935. 333.33e.30
768.3
DMpB
201.
206
473.4g22.)
836.9
543.7(
246.4m
771. (71.1
607.3):122
GOALFB
301.4:s.5
205.7
8711
887 2115))
776.6:44)
465. Sera:
9 19.88
768.6rs.a
Expert
910.9m..m
627.901.0
992.3ko.)
989.9
804.1q
934.0
971.6:1.a
985.7(1-4
Table 7: Cumulative reward for each task in the cheetah and quadruped environments, averaged over
repetitions. Experiments are done with K = 1 expert demonstrations. Standard deviation is reported in parenthesis.
E.3 FB MODEL QUALITY
In our experiments averaged the performance over 10 FB models. To evaluate the robustness of the pre-trained steps we report in Fig. the fraction of models F(T) with a combined normalized score above certain threshold Let M 10 be the number of FB models and x... be the score of model m
averaged over environments, tasks, FB-II. methods, and repetitions. then F(r) 1 ΣM Il [Imm >
32

---

Published as a conference paper at ICLR 2024
Algorithm
Time
DWBC
9h32m
Algorithm
Time
ORIL.
7h45m
Algorichm
Time
BCen
Im
ORIL-O
6h56m
BCDIAYN
ER
<5s
ORIL-OO
7h12m
ERDIAYN
<38
RERFB
<šs
OTRTDA
GOALDIAYN
BBELL
SQIL
10h18m 8m
GOAL-TD3
<5
FMpB
3m
MASKDP
FBLoS$pB
DEMODICE
12h59m
GOAL GPT
<5s
DMPB
3m
SMODICE
6h35m
GOALER
<5s
LOBSDICE
6h28m
IQLEARN
13h17m
Table 8: Average time for generating imitation learning policy once provided expert demon- strations. This clearly shows the advantage of BFMs that need only to infer one or multiple policies without explicit training.
will wore can
willuut warn stat
70
Inttatte Soore (5)
40
(a)
Bithwt
(b) ERFB
(c)
BCrawith
warm start
start
averaged over domains, tasks and repetitions for ERFE BCFB with and without warm start, using K = 1 expert demon- strations.
Figure 7: Imitation
Figure 6: Examples of trajectories in the maze with different FB- II. methods. The expert trajectory is reported in red while the II. trajectory is reported in green.
7] (Agarwal et al., 2021). We can notice that all the drop in probability is very concentrated towards high values (between 0.5 and 0.7) denoting a small variability in the performance of the FB models.
1.00
F(t): fractin of FB models with average reward >
0.75
0.50
Cheetah
0.25
Maze
Quadruped
Walke Walker
0.00
0.1
0.5 0.6
Average Reward (T)
Figure 8: Performance profiles on combined tasks and FB-IL methods. Shaded regions show point-
wise 95% confidence bands based on percentile bootstrap with stratified sampling.
E.4 EFFECT THE NUMBER DEMONSTRATIONS
We investigate whether the performance of the IL methods can improved by increasing the number of expert demonstrations. In Fig. 9 report the aggregate performance for all the tested algorithms.
33

---

Published as a conference paper at ICLR 2024
As expected, BC (similarly DWBC) improves significantly with the number of experts trajectories. However, even with 100 trajectories (corresponding to 1M number of expert transitions), BC only marginally outperforms FB-IL methods using only trajectory.
Similarly, DEMODICE and IQ-LEARN enjoy a significant improvement but they are still unable to match the performance of FB-IL methods.
Other baselines are practically unaffected by the number of expert trajectories given that the bottleneck is mostly the inefficient usage of unsupervised samples. FB-IL methods are overall stable achieving very good results already with a single trajectory. We think that increasing the number of trajectories does not help pre-trained models since are limited by the approximation errors of the model. We think we may be able to overcome this limit by allowing fine-tuning the model, but this is outside the scope of the paper.
Notice that many approaches GOALDIAY,, MASKDP, GOAL-TD3, GOAL-GPT, GOALFB- are not able to deal with multiple expert trajectories.
E.5 GENERALIZATION TO DIFFERENT STARTING POINT DISTRIBUTIONS
We investigate how well imitation algorithms generalize when the demonstrations are collected under a different initial state distribution than the one used for evaluation.
Specifically, we repeat the experiments of Section 5 using 200 new expert trajectories of 1000 steps generated by the same TD3-based expert policies of Section 5 but changing the initial state distributions as follows:
for the walker, quadruped, and cheetah domain, we initialize each trajectory directly in the expected long-term stationary behavior (i.e., walking position for the walk task, running one for run, etc.). Concretely, we achieved that by taking, for each domain and task, the state observed at step 500 of each of the 200 trajectories used in Section 5 (which made sure to be representative the desired stationary behavior) and by rolling out 1000 steps using the expert policy starting from each of these 200 new initial states.
for the maze domain, we randomly initialize the agent position in the upper-right room instead of the original upper-left one.
Then, we use these new demonstratinos to test all FB-II. methods as well as the most performing baselines from each II. algorithm family. Everything else (protocol, hyperparameters, etc.) is exactly the one used in Section 5.
In particular, each imitation policy is tested on the original initial state distributions of each domain (i.e., a random position and orientation for walker, cheetah, and quadruped, and a random position in the upper left room for maze), which are now very different from the one used to generate the expert demonstrations. The main challenge is that now the demonstrations are only showing the desired behavior (c.g.. how walk) but not how to reach that behavior (e.g., how move to a walking position when lying on the floor). The results are shown in Figure 10 and Figure 11.
In Figure 10, we report the average performance loss that each IL method suffers due to distribution shift (computed as the ratio between the imitation score obtained using the modified demonstrations this sections, and the one using the original demonstrations of Section 5).
Overall, methods in the BC and FM/DM classes seem to suffer the largest performance loss. Reward- based and goal-based methods are the least impacted by distribution shift, with the former being almost unaffected. Overall all the FB-IL methods consistently achieve good performance even under distribution shift (with a performance drop between 2% and 22% depending on the method), which is not the case for some baselines (e.g., those in the BC and FM/DM).
While goal-baue based methods slightly outperform the task-specific TD3IL baseline in our base setup (Fig. 1), this is not the case anymore in the generalization setup, with the best pretrained method now around 93% of the top non nretrrntd performance (Fig. 10). Still, the overall picture in Figs. I and 10 is largely the same.
34

---

Published as a conference paper at ICLR 2024
GRUL.C
col caz--00
OTNe
Toun
MANIDD
KOR ALT
RERRO
FRL FOLADdee
DMrs
Figure 9: Aggregate normalized score II. algorithms with different numbers expert demonstra- tions. Results are averaged over domains, tasks and repetitions.
E.6 DISTRIBUTION MATCHING: MATCHING THE SUCCESSOR MEASURE ON AVERAGE VS AT EACH STATE
The classical distribution matching loss (10) amounts to estimating the divergence between the overall occupation measures of a policy and the expert policy,
L (z) B R (so) Exprro (xo))
(94)
as Exp-po A is the cumulated discounted measure over all states visited by a when starting at 80 Pn-
35

---

Published as a conference paper at ICLR 2024
Olfect
SPM
PS-L (ein)
N1.21
$L.E
3608
TD:
PBG COLOONN
DM
Goole
RAN
FB4L
130
Looy Unnoorert Stondee Precosl
TDst
Duobice
GRAL- 7DA
RIW
Goncen
Figure 10: Generalization experiment: (top) average imitation score: (bottom) ratio between the average imitation with distribution shift (i.e., using the demonstrations with modified initial state distribution, as described in the first paragraph of Appendix E.5) and without (i.e., using the demonstrations of Section 5). Imitation scores are computed by averaging over all domains, tasks, and repetitions for single expert demonstrations.
DMC Charsa
DMC Mar
312
DMC Ouddaped
DNC Welker
NO.
Figure 11: Generalization experiment: imitation score for each domain averaged over tasks and repetitions for a single expert demonstration. The imitation score is the ratio between the cumulative return of the algorithm and the cumulative return of the expert.
But as mentioned in Section 4.3, we have chosen the loss
LR (2) Eggeps ||M (8O) M =m
(95)
36

---

Published as a conference paper at ICLR 2024
wann sirt
wilhut warm start
70
Imitatin Soow (%)
00
40
30
20
|HP
|B,
Figure 12: Difference between matching the successor measure at each state, i.e., using the loss Bearo. [| - = (s and matching it on average w.r.t. the imitial state, i.c., using the loss [B M= (so) Esa M (so)ll)) or Begr=p M (su) Esgrpo M M**(su) li...
instead. This compares the successor measures of the expert and π. each point separately.
This is a stricter criterion. For instance, take S' to be cycle S = {l, nt.. and take an expert policy that moves to the right on the cycle, 8 - measure Ex M (sa) of that policy is uniform too, and is the same as for policy that stays in place.
mod n. If Po is uniform, the overall occupation
Putting E., outside the norm is a simple way to solve this problem. Another way would be to consider distributions over state transitions (st: 8t+1) instead just states 8t, as done, e.g., in (Zhu al., 2020; Kim et al., 2022a), but this requires changing the foundation model.
In Figure 12, we report the difference overall performance between three such variants: putting Begrop inside the norm in classical distribution matching, putting Bagrp, inside the norm (thus widening the initial states to all states from which we can estimate successor measures from the demos), and putting Exgrop, outside the norm (our main choice). The norm chosen is Il-
The variant with Espop, inside underperforms, while the variant with Esoxpo inside only works well in the presence of our warm start initialization 20 + Ep, (B].
This is not surprising: indeed, with Exg~po inside the norm, the warm start is the only way the algorithm can incorporate information from the whole expert trajectory. Without warm start, it only gets information from the earliest part of the trajectory. More precisely, with the expectation E Esprpp inside the norm, the loss in Thm. becomes
LB (z)
Eag-ro [Cov B)F 2) -
B(sK+1) 18, Re Te]]1 est
(96)
from which it is clear that:
1. The model F only matters through F(so,
Along the expert trajectories, states far from the initial state distribution Po are discounted. So the later sections of the expert trajectories are largely discarded.
On the other hand, the version with Esoo-p outside does not rely on the warm start to incorporate information from the whole expert trajectory: it demonstrates greater robustness the initialization of z, for negligible difference in performance compared to -xo~,, fmside-with-warm--e corresponds to finer theoretical criterion for identifying policies, explained above.
It also
F WAYPOINT IMITATION LEARNING EXPERIMENT
We consider imitating sequence of yoga poses from the Yoga benchmark (Mendonca et al., 2021). The Robo Yoga benchmark defines 12 different robot positions for the walker and quadruped domains of the DeepMind Control Suite. Here we focus walker and generate expert demonstrations by concatenating different poses, keeping each fixed for 100 time steps. Figure 13 shows an example of such demonstration.
37

---

Published as a conference paper at ICLR 2024
Note that the expert is both non-stationary, the underlying task (i.e., reaching the pose) changes time, and non-realizable, as the transitions between poses are instantancous and thus not physically attainable. For this reason, compare only goal-based IL methods (GOALFB. GOAL-TD3, and GOAL-GPT) which are the only ones capable of dealing with non-stationarity.
I
Figure 13:
Example of sequence of yoga poses to imitate. From left to right: lie front, legs up, stand up, lie back, kneel, head stand, bridge.
F.1 DETAILED PROTOCOL
We evaluate each goal- baald IL method on 1000 sequences of poses. Each sequence is built by first sampling 10 out of the 12 yoga poses without replacement and then building a state trajectory of 1000 steps where each of the 10 poses is kept fixed for 100 consecutive steps. For each of the 1000 test sequences, each method is tested starting from a state randomly generated from the standard initial state distribution the walker domain.
Each generated trajectory is scored in terms of the rewards of the pose sequence it intends to imitate. More precisely, each pose in the Robo Yoga benchmark is associated reward function
{1
if St41 ollco I 0.55, otherwise,
($(+1)
where 9 is the state corresponding to the pose. Letting (gt)<2 denote a sequence of poses to imitate, the score we compute for the corresponding trajectory (s4)120 produced by II. method E gg++1 1s This number is then normalized by the total reward achieved by TD3 trained offline on each pose separately. Each number in Figure 4 is obtained by averaging the scores obtained by the II. method over the 1000 test sequences and pretrained BFMs, plus/minus the standard deviation of the average score of each of the 10 BFMs divided by V10.
F.2 ALGORITHMS AND HYPERPARAMETERS
GOALFB. We pretrain 10 FB models (with 10 different random seeds) as described in App. D.3 using the same hyperparameters reported in Table with only three modifications: we add extra hidden layer to the and actor networks, we reduce the learning rate of to 10-6 and we train for 3 x 10" gradient steps. At test time, we imitate pose sequences by using lookahead of 1. That is, at time of the produced trajectory we play = n2, (st) with (s ), where is the current state and sdoroo is the state one-step ahead in the demonstration (i.e., the next pose to imitate).
Goal-TD3 and Goal-GPT. We use the same 10 pretrained models considered in the experiments of Sect. 5.1 and 5.2. See App. C and for the detailed training protocol and hyperparameters. At test time, for GOAL-TD3 we use the same lookahead of I as for GOALPB- On the other hand, for GOAL-GPT we use lookahead of 16 as we found lookahead of to working poorly. This is likely due to the fact that GOAL-GPT is pretrained with context length of 16, and thus tries to reach a goal state 16 steps ahead with an autoregressive (i.e., history-dependent) policy. Setting the lookahead to essentially implies that we execute Markovian policy as the history is reset every single step.
F.3 QUALITATIVE EVALUATION
Figure 14 shows an example trajectory generated by GOALFR when imitating the pose sequence of Figure 13.
38

---

Published as a conference paper at ICLR 2024
H
Figure 14: Imitating the sequence of yoga poses from Figure 13 by GOALFB. The agent learns how quickly transition between each pose despite not being demonstrated how do
Besides being able to reproduce each of the seven poses, GOALFB capable of transitioning between them despite not being demonstrated how to do so. This is because the underlying FB model, from purely unsupervised pre-training, potentially learned how to reach each pose from any other state in the data.
39

---

Published as a conference paper at ICLR 2024
ORTLaG
OR
ORII--OOIQL
OTRiq
DEMODICE
SMODICE
LOBSDICE
halfchectah--medie cedirtaa-tmmmiinnncceere222
70.7a
83.0
78.5
70.3:2e)
(3.9
Booo;
69.9t..7
-
-ediun--ppla----22
35.0
34.2
35.1:0.5)
35.0)..)
32.0
310roa:
35.1
12.7
42.40.3)
42.6ru.)
42.5m.:)
42.6
42.5m..)
walker 2d medium--expet--22
103.112m
102.8;40)
105.4...
100.1:44:
99.8
107.91o..
106.2
walker2d medium medium replayy vv2
a2
43.4
42.1:2.0)
50.5.o.9)
36.3
20.2
37.4
walker2d 2--edd-----2222
61.7:24:
68.312.0)
62.1:2.5)
65.3/4.))
68.8:14)
53.0
56.9n1.)
Table Normalized score for imitation on the D4RL, benchmark with expert demonstrations.
G BASELINE RESULTS ON D4RL
In this section, we check that our baseline implementations achieve comparable performance to the results reported in the literature.
For this, we report the evaluation of few baselines on the standard D4RL benchmark (Fu et al., 2020). We consider the same setting as in (Luo et al., 2023). Similarly (Luo al., 2023), use IQL as the learning algorithm for ORIL and OTR, which we noticed to be performing better than TD3 in this particular benchmark.
Tab. 9 shows that our implementations achieve comparable performance to the results in the literature, confirming their correctness." This also shows that our setting particularly challenging for *DICE algorithms that are intrinsically tight to conservative updates.
II is likely that better results can he obtained through hyper-parameter tuning (which we did not perform in these experiments).
40


Reviews:
[{"review_id": "ZIDNPGOPqZ", "reviewer": "Reviewer_TWkJ", "summary": "This paper leverages advances in behavior foundation models (BFMs) based on successor measures to build a BFM that can be pre-trained on a dataset of unsupervised trajectories without any prior knowledge of subsequent imitation behaviors. To that end, they describe several imitation learning methods that are based on a pre-trained Forward-Backward (FB) model. In the experiments the authors verify if an FB model that is pre-trained on one environment is able to imitate a wide range of tasks given only a few additional demonstrations on the target task.", "questions": "Re future directions: Do the authors have any comments on utilizing FB pre-training for domain adaptation from sim to real? For example adapting to a locomotion task on a real robot from sim pre-training to overcome a low-dim parameter shift, and if this formulation will remain viable at all.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1) The paper is easy to follow and well-structured with the main contributions listed clearly in the introduction alongwith grounding in related works and building blocks that make up this method. The paper is also technically sound with no methodical kinks that I could find. \n2) Experimental details are sufficiently described for reproducibility.\n3) Experimental evaluation in the paper is very extensive covering a variety of methods with different pre-training methodologies on multiple tasks in the DeepMind Control suite.", "weaknesses": "Clarity: I think the clarity of the preliminaries can be improved. The paper assumes understanding of the FB framework which can make it hard to understand the core foundation of this work without going through previous papers. Effort should be made to make this work as self-contained as possible.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "9Vdhk7psvR", "reviewer": "Reviewer_bPor", "summary": "This paper leverages recent work on the successor measure, i.e., a generalization of the successor representation to continuous (or large) state spaces. In particular, the authors leverage the Forward-Backward (FB) representation, which learns a low-rank decomposition of the successor measure conditioned on a policy embedding z. These policy embeddings should be approximately optimal for a particular reward function. The authors show how to leverage the FB representation to perform various forms of imitation learning by reusing the FB representation and learning an appropriate policy embedding for the imitation task.\n\nThe work is well-motivated, and the importance is perhaps understated; imitation learning is a staple in many communities, and this paper provides notable improvements over well-established methods. As I see it, the primary contributions are:\n* The authors provide a novel application of the FB factorization of the successor measure to imitation learning.\n* Provide an extensive empirical analysis of FB applied to different forms of imitation learning, i.e., behavior cloning, reward-based imitation learning, distribution and feature matching, and goal-based imitation.", "questions": "* The DMC Maze tasks seem relatively straightforward; can you help give a better intuition as to why some methods fail to learn in this environment?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* The paper is very well written and easy to follow; each subsection detailing the application of FB is succinct but still provides many insights on the FB representation that prior work hadn’t touched on.\n* FB is unique in that it can be pre-trained from offline non-expert transitions and subsequently used to derive a reusable imitation policy from expert trajectories with very little overhead.\n* Well-designed empirical methodology with abundant baselines (both published and unpublished) on various continuous control domains. The experimental results are convincing and show significant improvement over the baselines.\n\nI also want to mention that I appreciate the level of detail in the supplementary materials and the lengths the authors went to describe how the baselines were implemented and tuned. This gives me confidence in the results and the ability to reproduce said results.", "weaknesses": "* I appreciate 21 tasks across four domains, but one concern is that the general findings might not extrapolate to other continuous control domains.\n* The results are limited to lower dimensional state spaces. I would appreciate at least discussing the method's limitations when scaling to higher-dimensional state spaces. It seems like FB has never been scaled beyond these low-dimensional continuous control tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "SnKMZF7sMf", "reviewer": "Reviewer_YAAN", "summary": "This work leverages recent behavior foundation models (BFMs) that are based on successor features for fast imitation without RL or fine-tuning. Notably, this work uses the forward-backward (FB) framework to build BFMs that can solve any imitation learning tasks. The resulting set of algorithms is called FB-IL.  FB-IL accommodates several IL principles, including behavior cloning, feature matching, reward-based, and goal-based reductions. Experiments show that FB-IL can produce imitation policies from new demonstrations within seconds.", "questions": "1.\tThis work leveraged BFMs based on successor measures, notably the forward-backward (FB) framework to solve fast imitation tasks. However, the introduction of successor measures is not clear enough. From section 3, we know that the successor measure describes the cumulated discounted time spent at each state following policy. The FB framework learns a tractable representation of successor measures that provides approximate optimal policies for any reward. What is the merit of achieving this goal?  And how does this property relate to FB-IL?\n\n2. Will the proposed methods require RL routine or fine-tuning in environments with high-dimensional states?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "It is interesting and novel to leverage behavior foundation models based on successor features to enable fast imitation learning. Building on top of the forward-backward (FB) framework, FB-IL has several important and useful properties:\n\n1.\tPre-training the BFM only requires a dataset of unsupervised transitions/trajectories.\n\n2.\tThe BFMs solve imitation tasks without any complex reinforcement learning problem, and the computation is minimal,\n\n3.\tThe BFM is compatible with different imitation learning settings, which is impressive. \n\n4.\tThe proposed FB-IL algorithms perform three orders of magnitude faster than offline IL methods, because FB-IL does not need to run full RL routines to compute imitation policies. \n\n5. It is nice to see the proposed framework can be extended to IL tasks with non-stationary demonstrations.", "weaknesses": "1.\tThe proposed method is motivated and based on the previous work in zero-shot reinforcement learning (Touati et al., 2023). However, the forward-backward framework contents in section 3 are not easy to follow. Don’t assume readers are familiar with the framework proposed by Touati et al.\n\n2. The proposed methods are only evaluated in the tasks with low-dimensional states. I am wondering if the proposed methods still do not require RL routine or fine-tuning in environments with high-dimensional states.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "svDGteELkD", "reviewer": "Reviewer_EALv", "summary": "The study extends prior research research into the Forward-Backward (FB) representation framework, a framework for training behaviour foundational models that hinges on representing agents + environment using an embedding space of policy representations and a set of two functions (F and B) that can be used to generate policy representations that are optimal with respect to emperically measuires rewards or reward functions unknown at model learning time.\n\nIn this study, the framework is extended to address Imitation Learning (IL), where policies need to be learned from one or more expert demonstrations. The paper introduces multiple FB-IL methods, each mimicking existing IL approaches. However, these FB-IL models, are trained using pre-trained domain-specific FB models (that were not exposed to reward or demonstrations of tartidular tasks), reducing computation time significantly. This is because he model of the environment and how the agent can interact with it are pre-trained, and IL training primarily focuses on finding the right policy embedding (z).", "questions": "I found it quite surprizing that the method cannot improve generated policies in any significant way with more than a single IL example. I wonder if that is a reflection of the simplicity of the environments (given that it *is* doing about as well as the expert that generated the single example)?\n\nIt was not clear to me how $\\pi_𝑧(𝑠)= \\arg \\max _a  F (s, a, z)^T z$ is implemented. Is the arg max optimization occuring at every step or is the policy represented in yet another model that needs to be trained to estimate this arg max? If so how fast is that trainig?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The various proposed methods are extremely sample efficient (doing about as well with a single expert example as they do with 200), are orders of magnitude faster to compute (excluding the FB training time which is performed once per domain and can then be used learn any task within that domain, so long as the space was adequarely explored at FB training time).\nThe results show that the IL-trained models can nearly match the performance of the expert models which they need to imitate (which were fully trained with RL).\n\nThere are many IL methods and this paper aims to show that the FB framework can support them all. As such, the paper lists the main IL approaches and for each, introduces an FB-based variant that mimics that approach. In that respect this paper presents not one but **multiple** methods that are quite distict from one another yet all sare the same (FB) representation framework, thus driving home the point that the FB framework is quite versatile and does not limit the IL approaches.", "weaknesses": "The basic FB framework is skimmed though very quickly, requiring the unfamiliar reader to go back to 1-2 previous papers on the topic to fully inderstand the representational framework.\n\nThe sheer number of \"new\" methods is a bit more that the reader might have bargined for in a 9-page paper (the very lengthy appendix serves as emphasis). Note that this can be seen as a strength of the work it self.\n\nPrior work on the FB framework has been limited to the somewhat simplistic gym environments, This work is no different. While the FB framework can scale by increasing the complexity of the embedding functions and the dimensionality of the policy representation space, it is unclear, at this point, that the framework can serve well in non-simulated environments. IL seems to me to be a practical method for real world problems and it would be nice to see FB applied to non-simulation (or more complex) environments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (spotlight)