[Imitation Learning] BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning

https://arxiv.org/pdf/2202.02005

https://sites.google.com/view/bc-z/home

bc-z

BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning

Introduction

로봇이 새로운 vision-based manipulation task에 대해 zero-shot or few-shot으로 generalization하기는 어렵다. 이 논문에서는 이 문제를 data collection을 위한 shared autonomy와 유연한 task conditioning을 통한 imitation learning으로 해결하고자 한다. 또한 이 논문에서는 training 과정에서 등장하지 않은 task에 대한 zero-shot과 few-shot generalization 등을 포함하여 다양한 task를 해결하는 interactive imitation learning에 대해 방대한 규모의 emprical study를 제공한다.

 

Related Work

기존의 imitation learning은 어느정도 generalization을 할 수 있다는 것을 보여줬다. 하지만 기존의 DAgger와 같은 data collection method는 효율성 문제가 있었는데, 최근 HG-DAgger와 같은 method는 이러한 효율성 문제를 해결하였다. 이 논문에서는 이 방법을 채택했다고 한다.

간단하게 설명하자면, DAgger(Dataset Aggregation)는 모델이 오작동하는 상황에서의 expert demonstration(전문가의 시연)을 통한 data collection 방법이다. 이때 모든 실패 상황에 대해 똑같이 데이터를 수집하기 때문에 많은 양의 데이터를 수집해야 해서 효율성이 떨어지는데, HG-DAgger는 이를 보완하였다. 모델의 confidence를 명시적으로 평가하여, 불필요한 데이터 수집을 최소화함으로써 효율성 문제를 해결했다.

이 논문에서는 처음보는 task에 대해 zero-shot or few-shot generalization을 목표로 한다.

 

Problem Setup and Method Overview

Policy architecture는 encoder, control layer 두 단계로 이루어진다. encoder 단계에서는 command(or human demo video)를 입력으로 받아 embedding을 출력한다. control layer에서는 RGB image와 encoder에서 나온 embedding을 입력으로 받아 action을 출력한다.

 

Data Collection and Workflow

데이터셋 수집 방법은 expert demonstration과 shared autonomy를 사용하였다. expert demonstration은 HG-DAgger에서와 같이 전문가의 시연 영상을 training할 때 데이터셋으로 사용하고, shared autonomy는 컨트롤러로 로봇의 움직임을 직접 제어한다는 개념으로 받아들이면 된다.

 

Learning Algorithm

Task embedding을 위한 encoder(아래 사다리꼴 모양)부터 살펴보자. 앞서 말했듯 이들은 language command 혹은 human demo video를 입력으로 받아 embedding을 출력한다. language encoder에 대해서는 pre-trained language embedding을 사용하고, video encoder에 대해서는 ResNet-18 based CNN을 사용하고 처음부터 끝까지 학습시킨다. video encoder를 학습시킬 때 loss 함수는 기존 behavior cloning의 loss 함수에 video encoder와 language encoder으로부터 나오는 embedding의 cosine distance를 더한 형태다.

Policy training은 RGB image를 data augmentation 과정을 거친 뒤 embedding하고, 신경망을 거쳐 action을 출력하도록 한다. XYZ와 axis-angle prediction에 대해서는 huber loss를 사용하고, gripper angle에 대해서는 log loss를 사용한다.

FiLM-conditioned task embedding은 입력 feature(vision network)에 conditional information(video or language embedding)을 적용하여 flexible하게 특성을 변환한다.

 

Experiment Results

Single-task imitation learning에서 좋은 성능을 보였다. 하지만 새로운 task에 대해서는 task마다 성능이 상이했다. 또한 ablation study로 control layer와 encoder 중 어느 부분이 bottleneck인지 확인한 결과, control layer 부분이 bottleneck임을 확인했다. embedding을 생성할 때, language command 보다 video로부터 생성하는 것이 더 어려웠다. 본문에서는 언급하지 않았지만, 논문에서 사용한 핵심 기술인 adaptive state diff scheme이 중요했다고 한다. 마지막으로 HG-DAgger로 data collection 한 것이 성능을 유의미하게 향상시켰다고 한다.

 

Discussion

이 논문의 핵심을 요약하면, (1) 간단한 imitation learning으로 새로운 task에 대해 generalization이 가능한 것을 보여줬다. (2) 또한, HG-DAgger로 data collection하는 것이 성능 향상에 큰 영향을 미친 것을 확인했다. (3) 마지막으로, language command embedding에서 pre-trained model을 사용해도 충분하다는 것을 보여주었다.

하지만 한계점도 존재한다. (1) 새로운 task에 대한 성능이 task마다 상이했다는 점과, (2) language command structure가 제한적이었다. (3) 또한 본문에서도 언급했듯이 video conditioning에 대해서도 큰 어려움이 있었다.