Paper: Fast Imitation via Behavior Foundation Models (ICLR-2024)

Reviewer comment (Writing):
The 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.

Author response:
-  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.

Reviewer comment (Experiments):
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.

Author response:
-   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.

Reviewer comment (Writing):
This 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?

Author response:
- 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)).

Reviewer comment (Reproducibility):
Will the proposed methods require RL routine or fine-tuning in environments with high-dimensional states?

Author response:
-   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.

Reviewer comment (Experiments):
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.

Author response:
- 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.

Reviewer comment (Evaluation):
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?

Author response:
- 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.