오늘은 「Where are we in the search for an Artificial Visual Cortex for Embodied Intelligence?」 논문을 리뷰해보려고 한다. 이 논문은 사전학습된 vision encoder를 기반으로 로보틱스 다운스트림 태스크를 해결하는 현재의 연구 흐름 속에서, 더 나은 모델을 설계하기 위해 우리가 반드시 짚고 넘어가야 할 핵심적인 내용들을 체계적으로 정리한다.
구체적으로는 vision encoder를 사전학습할 때 어떤 objective를 가진 모델이 유리한지, 데이터셋은 어떻게 구성하는 것이 효과적인지와 같은 질문들에 대해 다룬다. 또한 다운스트림 단계에서는 총 17개의 로보틱스 다운스트림 태스크를 선별하여 소개하고, 각 태스크에서 사전학습된 vision encoder를 기반으로 일반적으로 어떤 학습 방식이 사용되는지를 설명한다. 이를 통해 태스크의 특성에 따라 full fine-tuning이 적절한지, 혹은 linear probing이 더 효과적인지 등을 판단할 수 있는 기준을 제공한다.
따라서 이 논문을 읽고 나면, vision encoder를 사전학습할 때 어떤 요소들을 고려해야 하는지, 그리고 사전학습된 vision encoder를 로보틱스 다운스트림 태스크에 적용할 때 보통 어떤 학습 파이프라인을 따르는지에 대한 전반적인 그림을 이해할 수 있다.
https://arxiv.org/abs/2303.18240
Where are we in the search for an Artificial Visual Cortex for Embodied Intelligence?
We present the largest and most comprehensive empirical study of pre-trained visual representations (PVRs) or visual 'foundation models' for Embodied AI. First, we curate CortexBench, consisting of 17 different tasks spanning locomotion, navigation, dexter
arxiv.org
Writing flow
앞으로 전개될 글의 흐름은 다음과 같다.
먼저, 로보틱스 다운스트림 벤치마크인 CortexBench를 소개하고, 이를 통해 기존에 제안된 vision encoder 모델들을 공정하게 비교할 수 있는 기준을 마련한다. 이후, 여러 기존 vision encoder 모델들을 CortexBench 상에서 linear probing 방식으로 비교하며 각 모델의 성능을 살펴본다.
다음으로, CortexBench에 가장 유리한 vision encoder pretraining 데이터셋을 구성하고, 이 데이터셋에서 MAE Objective로 ViT-L를 학습시켜 CortexBench에서 SOTA 성능을 달성한 모델을 소개한다.
마지막으로, 각 로보틱스 다운스트림 태스크에서 (vision encoder - policy network)로 구성된 모델을 훈련시키기 위한 과정을 소개하며, 훈련시에 linear probing뿐만 아닌 다양한 fine-tuning 전략들을 정리하고, 태스크별로 어떤 fine-tuning 방법이 가장 효과적인지에 대해 논의한다.
Cortex Bench
CortexBench는 Embodied AI 커뮤니티에서 중요하게 여겨지는 17개의 로보틱스 다운스트림 태스크를 선별하여 구성한 벤치마크이다. 이 태스크들은 총 7개의 널리 사용되는 기존 벤치마크에서 가져온 것으로, 다양한 로보틱스 시나리오를 포괄하도록 설계되었다. 구성된 태스크들은 복잡한 manipulation 문제부터 locomotion, navigation에 이르기까지 폭넓은 범위를 포함하며, 사전학습된 vision encoder의 성능을 다양한 조건에서 공정하게 평가할 수 있도록 한다.
Adroit
Adroit은 28자유도를 가진 의인화된 로봇 손을 사용해 복잡한 manipulation 태스크를 수행하는 벤치마크이다. 대표적인 예로는 펜을 목표 위치로 옮기고, 동시에 정확한 orientation을 맞추는 태스크 등이 있다. 정교한 손 조작 능력을 요구하는 것이 특징이다.
MetaWorld
MetaWorld는 table-top 환경에서 다양한 manipulation 태스크를 수행하는 벤치마크이다. button-press, drawer-open과 같은 비교적 단순하지만 다양한 상호작용 태스크들이 포함되어 있으며, 로봇 조작의 범용성을 평가하는 데 자주 사용된다.
DeepMind Control (DMC)
DMC는 locomotion 태스크와 object manipulation 태스크를 모두 포함하는 벤치마크이다. 연속 제어 환경에서의 안정적인 정책 학습과 시각 기반 제어 능력을 평가하는 데 주로 활용된다.
TriFinger (TF)
TriFinger는 3자유도를 가진 로봇 손을 사용하여 Push-Cube 태스크를 수행하는 벤치마크이다. 비교적 단순한 구조의 로봇을 사용하지만, 정밀한 제어와 시각 정보 활용이 요구된다.
Habitat
Habitat은 navigation 태스크를 중심으로 구성된 벤치마크이다. 대표적으로 ImageNav와 ObjectNav 태스크가 있으며, 논문에서는 이러한 태스크를 해결하기 위해서는 visual representation이 고수준의 semantic 정보를 학습해야 한다고 강조한다.
Habitat 2.0
Habitat 2.0은 7자유도 로봇 팔과 suction gripper를 사용해 실내(아파트) 환경에서 객체를 재배치하는 rearrangement 태스크를 다룬다. 대표적인 예로 Mobile Pick 태스크가 있으며, manipulation과 navigation이 결합된 보다 현실적인 시나리오를 포함한다.
Performance Comparison of existing Vision Encoders on Cortex bench
위의 표를 통해, 기존에 제안된 여러 vision encoder들을 CortexBench 태스크들에서 비교한 결과를 확인할 수 있다. 모든 실험은 linear probing 방식으로 진행되었다.
실험 결과를 보면, R3M은 Adroit, MetaWorld, DMControl 태스크에서 가장 우수한 성능을 보였고, MVP는 TriFinger, ImageNav, Mobile Pick 태스크에서 최고 성능을 기록했다. 한편, CLIP은 ObjectNav 태스크에서 가장 뛰어난 성능을 보였다.
이러한 결과는 vision encoder마다 강점을 보이는 태스크 유형이 다르며, 사전학습 방식과 데이터셋의 특성이 다운스트림 태스크 성능에 큰 영향을 미친다는 점을 보여준다.
Scaling : vision encoder pretraining (Making a Good Vision Encoder Pretraining Dataset for CortexBench Downstream Tasks)
논문에서는 vision encoder를 사전학습(pretraining)할 때 scaling이 로보틱스 다운스트림 태스크 성능에 어떤 영향을 미치는지를 분석하기 위해 일련의 실험을 수행한다. 더 나아가, 만약 scaling된 데이터셋이 실제로 효과적이라면, 이를 활용해 학습한 모델이 CortexBench에서 SOTA 성능을 달성할 수 있는지를 검증하는 것이 목표이다.
Scaling 실험을 위한 사전학습 objective로는 MAE를 선택했으며, 모델 크기는 ViT-B와 ViT-L 두 가지를 사용했다. 데이터셋은 크기와 다양성의 정도에 따라 총 5가지 구성을 준비했다.
가장 기본이 되는 데이터셋은 Ego4D로, 일상적인 1인칭 시점에서 촬영된 비디오 데이터로 구성되어 있으며, manipulation과 관련된 장면이 비교적 많이 포함되어 있다. Ego4D+M은 Ego4D에 manipulation 중심의 데이터를 추가한 구성이고, Ego4D+N은 Ego4D에 navigation 관련 데이터를 추가한 것이다. Ego4D+MN은 manipulation과 navigation 데이터를 모두 포함한 버전이며, 마지막으로 Ego4D+MNI는 Ego4D+MN에 ImageNet을 추가하여 시각적 다양성을 극대화한 데이터셋이다.
먼저 모델 크기의 영향을 살펴보면, 사전학습된 vision encoder의 크기를 키울수록 대부분의 다운스트림 태스크에서 성능이 일관되게 향상되는 것을 확인할 수 있다.
다음으로 데이터셋 다양성이 미치는 영향을 보면, pretraining에 사용되는 데이터의 diversity가 증가할수록 downstream task들의 전반적인 성능이 크게 향상되는 경향이 나타난다. 다만 Ego4D 자체에 이미 manipulation 관련 장면이 상당수 포함되어 있기 때문에, Ego4D+N에서 Ego4D+MN으로 확장했을 때의 성능 증가는 비교적 미미한 편이다.
여기까지의 결과는 어찌 보면 직관적인 내용일 수 있다. 하지만 논문에서 특히 강조하는 중요한 포인트는, scaling은 평균적인 성능을 향상시키지만, 모든 개별 태스크의 성능을 항상 개선하지는 않는다는 점이다.
예를 들어 manipulation 다운스트림 태스크의 경우, vision encoder를 Ego4D+M으로 사전학습했을 때가 Ego4D+MNI로 학습했을 때보다 더 좋은 성능을 보이기도 한다. 이는 데이터 측면에서 볼 때, 다운스트림 태스크에서 자주 등장하는 장면을 사전학습 단계에서도 충분히 보게 하는 것이 중요함을 시사한다.
극단적으로 생각해보면, 평균 downstream success rate를 계산할 때 manipulation 태스크가 압도적으로 많다면, Ego4D+MNI보다 Ego4D+M이 더 좋은 선택일 수 있다. 그러나 CortexBench 전체를 기준으로 보면, manipulation 태스크와 navigation 태스크가 모두 포함되어 있기 때문에, 사전학습 단계에서 관련된 장면들을 폭넓고 균형 있게 관찰하는 것이 전반적인 성능 향상에 유리하다고 결론지을 수 있다.
VC-1
논문에서는 앞서 소개한 Ego4D + MNI 데이터셋을 사용해 MAE objective로 사전학습한 모델을 VC-1이라고 정의한다. VC-1의 다운스트림 성능을 기존의 여러 vision encoder들과 linear probing 방식으로 비교한 결과, CortexBench 전반에서 SOTA 성능을 달성했음을 확인할 수 있다.
이러한 결과는 MAE라는 학습 objective 자체의 효과도 크지만, 동시에 Ego4D + MNI 데이터셋의 구성과 다양성이 다운스트림 성능 향상에 결정적인 역할을 했음을 시사한다. 즉, 적절한 pretraining objective와 함께, 다운스트림 태스크 분포를 충분히 반영한 데이터셋 설계가 얼마나 중요한지를 잘 보여주는 사례라고 볼 수 있다.
How to train Downstream Tasks using pre-trained vision encoder
이번 session에서는 각 downstream task에서 (pre-trained vision encoder - policy network)로 구성된 모델을 어떻게 훈련시키는지 알아보고, 훈련 시에 linear probing 뿐만 아니라 다양한 fine tuning 기법들에 대해서 다룰 것이다.
먼저, 각 downstream task에서 (pre-trained vision encoder – policy network) 로 구성된 모델이 어떻게 훈련되는지 살펴본다.
Train Downstream Tasks : MuJoCo Tasks (Adroit, MetaWorld, DMC)
Adroit, MetaWorld, DMC benchmark에 포함된 MuJoCo 기반 task들은 모두 behavior cloning (BC) 방식으로 훈련된다. Behavior cloning은 사람이 미리 해당 환경에서 수행한 여러 개의 정답 trajectory, 즉 (state, action) sequence 를 수집한 뒤, 모델이 특정 state에 대해 동일한 action을 출력하도록 지도 학습하는 방법이다. 다시 말해, 모델은 전문가의 행동을 그대로 모방하도록 학습된다.MuJoCo task의 경우, 각 task당 제공되는 trajectory 수는 비교적 적으며, 약 25–100개의 sequence 만이 사용된다.
모델 구조를 조금 더 자세히 살펴보면 다음과 같다. 먼저 vision encoder 는 입력 이미지로부터 feature를 추출하며, 이 중 CLS token 만을 policy network의 입력으로 사용한다. Policy network는 3-layer MLP 로 구성되어 있으며, 시간적 정보를 반영하기 위해 frame stacking 기법을 사용한다. 즉, 최근 N개의 frame을 concatenation 하여 policy network의 입력으로 제공한다.
훈련은 총 100 epochs 동안 진행되며, 모든 epoch에 대해 evaluation을 수행한다. 이 중 test 성공률(success rate)이 가장 높았던 epoch의 결과를 최종 성능으로 report한다.
Train Downstream Tasks : TriFinger Tasks
다음으로는 TriFinger task에 대해 살펴본다. TriFinger task 역시 behavior cloning 방식으로 훈련되며, 각 task마다 100개의 trajectory가 제공된다.
모델 구조에 대해 살펴보면, Vision encoder에서는 CLS token을 추출하여 policy network의 입력으로 사용하며, policy network는 3-layer MLP로 구성된다.
훈련 설정은 task에 따라 다르다. Reach-Cube task는 100 epochs, Push-Cube task는 상대적으로 더 긴 1000 epochs 동안 훈련된다.
Train Downstream Tasks : Habitat Tasks (ObjectNav, ImageNav, Mobile-Pick)
다음으로는 Habitat task들에 대해 살펴본다. 먼저 ObjectNav는 약 77,000개의 trajectory를 사용하여 behavior cloning 방식으로 훈련된다. 반면, ImageNav와 Mobile-Pick은 총 5억(500M) step에 걸쳐 DD-PPO와 VER 기법을 적용한 순수 reinforcement learning 방식으로 학습된다.
이러한 Habitat task들은 공간적·시각적으로 매우 dense한 정보를 요구하기 때문에, policy network의 입력으로 CLS token이 아닌 patch token 전체를 사용한다. 이를 통해 장면 내의 세밀한 공간 구조와 객체 정보를 보다 풍부하게 전달할 수 있다.
또한 시간적 정보를 효과적으로 반영하기 위해, vision encoder와 policy network 사이에 2-layer LSTM과 GRU를 추가한다. 이 recurrent module들은 과거 frame에서 누적된 정보를 유지함으로써, policy network가 장기적인 시각–공간적 맥락을 기반으로 의사결정을 내릴 수 있도록 돕는다.
fine-tuning strategies for each task
이제 각 task별 훈련 과정에 대해서는 전반적인 감을 잡았을 것이다. 남은 핵심 질문은, 이러한 downstream task들을 학습할 때 vision encoder를 고정한 채 linear probing을 사용할 것인지, 아니면 vision encoder까지 포함해 전체 모델을 finetuning할 것인지에 관한 선택이다.
먼저 vision encoder를 얼리지 않고 downstream task에 성공적으로 적응시킬 경우, 어떤 이점을 얻을 수 있는지 살펴보자. 크게 두 가지 장점이 있다.
첫 번째는 feature extraction level이 task에 맞게 조정될 수 있다는 점이다. 어떤 task에서는 edge나 texture와 같은 low-level feature가 중요할 수 있고, 반대로 어떤 task에서는 object-level 또는 semantic-level의 high-level feature가 더 중요할 수 있다. Vision encoder를 함께 학습시키면, 이러한 feature hierarchy가 downstream task의 특성에 맞게 자연스럽게 재조정될 수 있다.
두 번째 장점은 domain gap을 완화할 수 있다는 점이다. 본 논문에서 다루는 downstream task들은 모두 simulation 환경에서 수행되지만, vision encoder의 pre-training은 실제 이미지 데이터셋을 기반으로 이루어진다. 이로 인해 필연적으로 sim-to-real domain gap이 발생하며, finetuning 과정을 통해 이러한 차이를 줄일 수 있다.
이러한 관점에서 논문은 vision encoder를 downstream task에 적응시키는 방법으로 두 가지 adaptation 전략을 제안한다.
첫 번째는 end-to-end finetuning 방식으로, vision encoder와 policy network를 포함한 모든 파라미터에 대해 gradient가 흐르도록 학습하는 방법이다. 이는 일반적으로 우리가 알고 있는 full finetuning에 해당하며, feature extraction adaptation과 domain gap 완화를 동시에 달성할 수 있다.
두 번째는 domain gap adaptation 방식이다. 이 방법에서는 먼저 pre-trained vision encoder만을 분리하여, downstream task의 scene 이미지들을 이용해 vision encoder를 사전 finetuning한다. 이후 실제 downstream task 학습 단계에서는 vision encoder를 고정한 채, linear probing 방식으로 policy network만 학습한다. 이 접근법은 feature hierarchy 자체를 바꾸기보다는, 주로 domain gap을 완화하는 데 초점을 둔다.
각 adaptation 방법의 성능을 비교해보면, downstream task 학습 시 충분한 데이터가 제공되는 경우와 데이터가 제한적인 경우 사이에 뚜렷한 차이가 나타난다.
먼저 ObjectNav, ImageNav, Mobile-Pick과 같이 대규모 데이터가 제공되는 task들의 경우, full finetuning이 가장 우수한 성능을 보인다. 이 경우 full finetuning은 앞서 언급한 두 가지 장점, 즉 task에 맞는 feature extraction level의 적응과 domain gap 완화를 모두 효과적으로 달성하며, 결과적으로 최고 성능을 기록한다.
반면, downstream task 학습에 사용되는 데이터가 부족한 다른 task들에서는 상황이 다르다. 이러한 경우 full finetuning은 모델이 새롭게 학습해야 할 정보에 비해 데이터가 지나치게 적어, overfitting이 발생하기 쉽다. 그 결과, full finetuning보다는 domain gap adaptation (MAE adaptation) 방식이 압도적으로 더 우수한 성능을 보이는 것을 확인할 수 있다.
새로 알게 된 것 : 기타
- 다른 pre-training 기법들에 비해서 MAE가 학습 비용이 낮은 편이다. 왜냐하면 encoder를 non-masked patch들만 통과하기 때문이다. Dino는 encoder를 두번 돌려야하고, 모든 patch를 넣어야 한다.
Comment
Vision for Robotics를 위한 기초적인 지식을 잘 쌓을 수 있는 논문이었다.