Fast Imitation via Behavior Foundation Models

We thank the reviewer for their comments and feedback. We are glad the reviewer appreciated the broad applicability of FB-IL across several IL paradigms.

• "Performance as a function of number of demonstrations". Most of experts tend to rapidly converge to a fixed stationary distribution (e.g., walking) independently from the initial state. Given that the dynamics is deterministic, imitating the long-term state visitation of a single trajectory is often sufficient to produce good behaviors. This is well illustrated in Sect. E.5, where even when the initial states used to demonstrate behaviors are different from the initial distribution used to test the imitation policy, the overall performance is still good. A second aspect is that FB-IL uses a frozen pre-trained model. As such, even with an infinite number of demonstrations we may not be able to perfectly represent the expert's policy due to the unavoidable approximation in the set of policies $(\pi_z)_z$, which are parametrized by the 50/100-dimensional vector $z$. On the other hand, BC performance keeps increasing with the number of demonstrations, as it fully exploits the capacity of the policy network, which has ~3K parameters. In order to match similar performance, we would need to fine-tune our model when a large number of demonstrations is available. Finding a proper fine-tuning procedure to well balance the efficiency of the pre-trained policies and the accuracy of an individually trained policy is an interesting venue for future research.
• Computing $\pi_z(s) = \arg\max F(s,a; z)^\top z$}. As the action space is continuous, we cannot directly compute the greedy action in each state. We use an policy network instead (see App. D.3 for details), which is trained to compute the greedy policy during the learning process of $F$ and $B$ following a TD3-like scheme. The overall structure is similar to the one used in (Touati et al, 2023). We will add more details in the final version of the paper.

We thank the reviewer for their comments and feedback.

• Thanks for your suggestion on improving the clarity of the preliminaries. We will revise the section to provide more insights on the FB framework, making sure the paper stands on its own more effectively.

• Q: "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?"

  The FB framework is designed to learn several components, such as a low-rank decomposition of the successor measure via $F$ and $B$, and a set of optimal policies $\pi_z$. Crucially, $F$ and $B$ are used to encode the optimal policies for any given reward (in the limit). At high level, the successor measure establishes a linear mapping between reward and the Q-function, expressed as $Q^\pi(s, a) = \int_{s'} M^\pi(s, a, ds) R(s')$. This property allows FB to recover the optimal policy for any reward specificed at test-time in a zero-shot manner, i.e., through a closed-form solution for the latent variable $z$.

  Different aspects of the FB model are crucial to derive different FB-IL methods, depending on the family of IL approach. For Behavior Cloning, we directly use the set of FB policies encoded by the latent variable $z$ and leverage the inductive bias of FB for learning optimal policies. For DM and FBLOSS methods, we directly leverage the fact that $F$ and $B$ approximate the successor measures. Finally, the zero-shot FB property of deriving an optimal policy from any reward function is directly used in reward-based IL. Indeed, consider the reward $r(s) = \frac{\rho^e(s)}{\rho(s)}$ used in SMODICE, or $r(s) = \frac{\rho^e(s)}{\rho^e(s) + \rho(s)}$ used in ORIL. Prior IL methods require training a discriminator to learn these reward functions and subsequently using an RL subroutine to maximize the reward, but FB-IL can instantly recover the optimal policy for these rewards by appropriately setting the variable $z$ (see Eq. (8) and Eq. (9)).

• Q: "Will the proposed methods require RL routine or fine-tuning in environments with high-dimensional states?"

  We focus on DMC continuous-control tasks to systematically evaluate various IL methods (both newly introduced methods and baselines) and this setting remains challenging for many state-of-the-art baselines, with only FB-IL methods consistently demonstrating both generality and strong performance. Furthermore, in the context of learning from proprioceptive states, quadruped is already relatively high-dimensional with 78-dimensional state and 12-dimensional action spaces. Nonetheless, we acknowledge that scaling FB approaches to image-based RL is an important direction for future investigation. Recent works such as [1] and [2] showed that low rank decompositions of dynamics and successor measures is indeed feasible even for problems with high-dimensional perceptual input, such as in ATARI games. This makes us confident that FB could perform well in such settings as well.

  Finally, while we agree that fine-tuning may be necessary for solving more complex tasks (e.g., reaching very narrow locations in a maze or controlling highly unstable robots), we believe that the behavior returned by FB would often offer a good initial solution, thereby significantly accelerating any subsequent finetuning.

Please let us know if there are any remaining questions or concerns.

References:

[1] Reinforcement learning from passive data via latent intentions, Dibya Ghosh, Chethan Bhateja, Sergey Levine, ICML 2023

[2] Proto-Value networks: Scaling Representation Learning with Auxiliary Tasks, Jesse Farebrother, Joshua Greaves, Rishabh Agarwal, Charline Le Lan, Ross Goroshin, Pablo Samuel Castro, Marc G Bellemare, ICLR 2023.

We thank the reviewer for their comments and feedback. We are glad that the reviewer appreciates our paper.

• Q: "Extrapolation to other continuous control domains?" While the domains considered in the paper may not be representative of all types of systems and tasks, we believe they already provide a good coverage of the main challenges arising in continuous control: Cheetah has very asymmetric dynamics where moving forward and backward is very different. Walker is an unstable system, as the agent would naturally tend to fall if not properly controlled. Quadruped is mostly stable, but it has high-dimensional state and action space. Maze is a stable environment (i.e., the agent remains in the same position if no action is taken) but it has a highly discontinuous dynamics close to the walls and the tasks we consider require a very fine control (e.g., to reach a specific location and speed). Given the current results on these domains, where our methods perform well and are very consistent (see e.g., Fig.8 for model robustness and the almost fixed set of parameters used across environments in Table 1), we are quite confident they could obtain good performance in other continuous control domains as well.

• Q: "Scaling to higher-dimensional state spaces?" The FB model is primarily designed to capture the dynamical aspects of the policy-environment interaction. In this sense, we believe that working on high-dimensional proprioceptive states (e.g., quadruped has 78d states and 12d action spaces) already challenges the core aspects of the model and allow for reliable comparisons with the state-of-the-art methods. On the other hand, learning from images may add a perceptual representation learning aspect that may bias the comparison with other baselines. Furthermore, recent works such as "Reinforcement learning from passive data via latent intentions" and "Proto-Value networks: Scaling Representation Learning with Auxiliary Tasks" showed that low rank decompositions of dynamics and successor measures is indeed feasible even problems with high-dimensional perceptual input, such as in ATARI games. While these algorithms still require fine-tuning at test time, we believe this suggests that scaling FB representation to image-based RL may be indeed feasible.

• Q: "Intuition of failures of some baselines". Some imitation learning baselines fail primarily because of issues with their offline training procedure or the overfitting of their reward learning. For instance, the family of *Dice algorithms regularize the occupancy state distribution induced by the learned policy to closely align with the training state distribution. *Dice algorithms demonstrated good performance in the D4RL benchmark (see appendix G), where datasets have a strong bias towards the specific task the agent is trying to imitate at test time. Indeed, we realized the balance between offline and expert data is crucial for *Dice algorithms to obtain a satisfactory performance. In our setting, RND datasets have a state distribution very different from any expert distribution and the conservative updates of *Dice algorithms often prevent them to properly leverage the expert data and learn a good imitation policy. In the specific case of the maze environment, we noticed that the discriminator $D(s,a)$ used in ORIL to define the reward function tends to overfit on the action component, while completely disregarding the state. This issue arises because, in goal-reaching tasks, most expert actions are zeros. In contrast, ORIL-O learns a discriminator that relies solely on the state, thus avoinding this problem and resulting in improved performance.

We spent a considerable amount of time tuning all baselines and overall we found that most of offline imitation learning algorithms are very sensitive to parameters such as expert-to-offline data ratio, offline RL subroutine used to optimize the imitation policy, and availability of state-action or state-only information. On the other hand, the "pre-train and then imitate" pipeline we employ for our model appears to more robust to the training data and imitation tasks.

We thank the reviewer for their comments and feedback. We are glad that the reviewer appreciates our paper.

• Thank you for your feedback on the clarity of the preliminaries. We will revise the section and the appendix to provide more details on the FB framework (e.g., how it is trained) to ensure that the paper is more self-contained.

• Q: "FB for sim-to-real". Thanks for suggesting the direction of applying FB models in sim-to-real scenarios. This is a very important problem, but unfortunately we believe that FB pre-trained models would not be robust to domain shift out-of-the-box. Although FB exhibits task generalization capabilities, it is crucial to recognize that the learned representations are closely tied to the dynamics (and samples) of the training environment. Nevertheless, we believe it is easy to integrate existing domain adaptation principles in the FB model to solve sim-to-real problems. The most immediate approach would be to fine-tune FB models with additional samples from the real environment. For instance, we could first use zero-shot properties of FB (either from reward or demonstration) to extract a policiy and then fine-tune it for the task at hand by using any standard RL algorithm. If additional unsupervised data from the real environment are available, we could fine-tune the forward representations $F$, which is the component most sensitive to the dynamics, while keeping the same backward embedding $B$ fixed, or we could fine-tune the whole model. If we have access to a parametrized form of the environment, we could meta-learn the FB model on randomized instances while conditioning FB representations on the history of past interactions (e.g., replacing MLP implementations of $F$ and $B$ with transformer architectures) to infer the right dynamics and thus improve adaptability.