Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes an imitation learning algorithm, TRAIL, that can benefit from a large amount of suboptimal demonstrations besides a small amount of high-quality demonstrations. This is achieved through learning a factored transition model with action reparameterization from the suboptimal or even random demonstrations before doing behavior cloning. The authors analyze the error bound of the algorithm and show improved sample complexity with action reparameterization. Moreover, TRAIL is verified on a set of navigation and locomotion imitation learning tasks and it outperforms baselines based on temporal action abstraction in terms of task success rate. 
main_review: ### Strength ###
- This paper considers a practical imitation learning setting in which there’s a very small amount of high-quality demonstrations that are expensive to obtain and there’s a large amount of (task-independent) suboptimal data that is cheaper to get. 
- The proposed algorithm is simple and theoretically-grounded. 
- This method of learning action reparameterization is applicable beyond imitation learning. Potential applications include offline reinforcement learning and transfer learning. 

### Suggestions and Questions ###
- In the analysis, the paper considers the discrete state (and discrete action) setting. Is it possible to extend the analysis to continuous spaces? 
  - In Theorem 3, when the latent action space is continuous, how does \pi_{alpha^*} in the second error term adapt from the first equation in page 5 which requires the latent action space to be discrete? 
  - Sample complexity is studied under the tabular case in page 5. What’s the sample complexity (e.g. Theorem 2) when latent action space is continuous? 
- In the TRAIL EBM for Theorem 1 paragraph in Section 4.3, you wrote “We allow \phi to be updated while optimizing J_{DE}”. Would this increase the first error term (Eq. (1))? Have you tried to optimize \phi, T, and \pi_\alpha simultaneously?
- I am excited about the experimental results. They look very promising. 
  - However, It’s not clear to me whether learning action reparameterization is superior to learning state encoding and learning transition models without factorization. For example, a naive baseline may be interesting to consider is learning a transition model from the large amount of suboptimal data and then applying GAIL or AIRL to match the state distribution of the small amount of demonstrations using the learned transition model. 
  - Some form of visualizing the latent action space learned may be interesting. 

### Other Comments ###
- \pi_k in the second equation in page 5, \pi_5 has not been introduced 
- In Theorem 1. Is \pi_Z S -> \Delta(Z)?
- In Section 4.3 TRAIL EBM for Theorem 1, “we set \rho to be the distribution of s’ in d^{off}”, does that mean you sample both the samples in the contrastive loss from the suboptimal data?

### Post-rebuttal comments ###
Thank the authors very much for their response in a short amount of time. All of my questions have been clarified. As before, I am leaning towards accepting the paper. 

summary_of_the_review: Because of the strengths discussed above, I am leaning towards accepting the paper.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: The paper considers an imitation learning (IL) problem with both expert and suboptimal demonstrations. The paper claims that sub-optimal demonstrations can be used to learn latent action abstractions which can improve the efficiency of down-stream IL. To solve this problem, the paper proposes TRAIL, which pre-trains an action encoder-decoder and a latent transition model using sub-optimal data, and performs behavioral cloning (BC) with expert data to learn a policy in the latent action space. The paper derives error bounds showing that the pre-training step and the down-stream BC step contribute to solving an IL problem in the original space from the view of divergence minimization. The paper also derives a bound showing that an optimal action abstraction may improve the sample complexity of BC. Experiments  indicate that TRAIL effectively learns from a limited number of expert demonstrations and is robust to the sub-optimality of sub-optimal demonstrations. 

Main contributions: An effective method to learn an expert policy from few expert data and a large set of sub-optimal data. The method is theoretically justified and empirically well-supported. 

main_review: ## Strengths

- The paper proposes a theoretically justified method to learn from expert and sub-optimal data. The method is quite flexible in the sense that it requires a minimal assumption on the sub-optimality of the data (i.e., only that sub-optimal data should have an overlapped support with the expert data). Though, the sample complexity result relies on the assumption of the optimal action encoder which may not be realizable in practice.

- The experiments are done thoroughly to evaluate the efficiently and robustness of TRAIL. A sufficient number of baselines is compared against TRAIL in a sufficient number of benchmark tasks. Still, the results could be made stronger by considering more practical tasks such as robot arm manipulation.  

## Weakness 

- The main issue is on the clarity of the paper. In particular, I do not understand $\theta_s$ and $\theta$ is section 4.2. Is $\theta_s$ the $s$-th column of the matrix $\theta \in \mathbb{R}^{d \times |S|}$? If so, is $\theta : S \mapsto \mathbb{R}^d$ a typo or you assume it is a linear map? I also do not fully understand about the linear transition model in section 4.3 and the cosine non-linearity. Why is the cosine applied here when linearity is the goal? 

## questions
- How are the number and dimensionality of latent action chosen for TRAIL EBM and TRAIL linear in each task? Are they tuned differently for each task or the same value is used across tasks?

summary_of_the_review: I find the problem setting and proposed method very interesting. The theoretical analyses are valuable, and the experiments are convincing. Though, there is an issue regarding the clarity which make it difficult to fully understand the paper. I rate the paper as above the borderline. 

** Update after authors' rebuttal **
The rebuttal and the revision address my concerns on the paper. I think more positive about the paper and have increased the score. The new result (Figure 7) in the appendix is a good addition and I suggest including it in the main paper as well.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes a method to accelerate behavioral cloning (BC) (especially in the low data regime) by utilizing a (much larger) auxiliary dataset of suboptimal behaviors. 
The authors claim that learning a good latent action representation (in this case, by learning a transition-based action representation from the auxiliary dataset) should make the imitation learning problem easier. 
The training is divided into two phases: (1) learning a factorized transition model in conjunction with an action embedding and action decoder, (2) learning a latent policy via the BC objective hoping that the transition-based representation accelerates the BC learning.
main_review: # Strengths:

(1) The method is tested in several challenging environments. In all of them, TRAIL performs better than or is on par with the baselines.  

(2) The paper is well written, and the theory appears to be sound

# Weaknesses: 

(1) The bound of the imitation learning error derived in theorem 1 and 2 relies on the fact that $|Z| < |A|$. In all of the experiments, the action embedding ($> 64$) is bigger than the original action space ($=8$ for the ant). Could the authors provide an intuition of why such big action embeddings are required? This is somewhat contradictory to the theory.

(2) The comparison to the offline reinforcement learning setting is interesting. 
However, the big promise of offline RL is to extract an optimal policy of the MDP induced by the offline dataset. That means, in principle, offline RL should be able to recover an optimal policy just from random data (given a sufficient state and action space coverage in the offline dataset) or a multitude of behaviors.
I would like to see a small paragraph in section 5.3 discussing this setting and how this affects BC. 

(3) In C.2, the authors write that the action decoder was fine-tuned on the expert data while keeping the action representations fixed. How does that affect the final performance?

# Needs clarification and general comments

(1) The authors claim that “[a] good action representation should enable good imitation learning.” 
Could the authors give an intuition (independent of the empirical results) on why a good BC action representation should emerge from a good transition-based representation?  

(2) In theorem 1 $\pi_{Z}$ is defined as $\pi_{Z} \colon S \times Z \to \Delta(A)$. If I am not mistaken, that should be $\pi_{Z} \colon S \to Z$

(3) I think in the sample complexity paragraph after (1) and (2) $J_{T}$ and $J_{DE}$ are missing, respectively.

(4) In equation 5 $\psi(s^{\prime})$ appears for the first time. However, I am failing the find the definition of $\psi(s)$

(5) The BC baseline is only trained with the expert data. Is there a way of incorporating the suboptimal data in the training of the vanilla BC baseline? Maybe by first training with the suboptimal data and later fine-tuning on the expert data.

(6) I would like to see a plot showing the effects of the different design choices (size of the latent embedding, fine-tuning the action encoder) on the performance.

(7) To empirically prove the theoretical claims, the authors could run another experiment in which they artificially blow up the action dimensionality and check if they can learn a much smaller embedding.

summary_of_the_review: At the moment, I tend towards a weak reject because I don’t find the connection between theory and empirical evaluation too convincing. The reduction in the BC error shown in theorem 1 and 2 are based on the assumption that $|Z| < |A|$. However, if I am not mistaken, that is not the case in any of the experiments where $|Z|$ is usually much bigger than $|A|$. 
On the other side, the empirical results show a clear improvement over the baselines and vanilla BC, indicating that the learned latent embedding is helpful. 
I am willing to change my score if the authors strengthen the connection between the theory and the empirical results. For instance, they could explain why the transition-based representation might help BC, although the latent embedding is bigger than the original action space.

### Decision due to rebuttal

The authors addressed all of my concerns sufficiently well. Therefore I am happy to raise my score from a 5 to a 6.
review_title: null

Meta-review: The paper investigates what we can learn from _suboptimal_ demonstrations for imitation learning. It suggests that we can learn about the structure of the environment by finding a factored dynamics model including a latent action space. It demonstrates both theoretically and empirically that this information can reduce sample requirements for downstream IL.

The reviewers praised the simplicity of the method (including its minimal assumptions), the theoretical analysis, and the breadth of the experimental validation. The authors were helpful during the discussion period, and addressed any questions or concerns the reviewers raised.

Overall, this is an interesting idea and a well-executed paper.
Decision: Accept (Poster)