Theia: Destilando Modelos Fundamentales de Visión Diversos para el Aprendizaje de Robots

Theia: Distilling Diverse Vision Foundation Models for Robot Learning

July 29, 2024

Autores: Jinghuan Shang, Karl Schmeckpeper, Brandon B. May, Maria Vittoria Minniti, Tarik Kelestemur, David Watkins, Laura Herlant

cs.AI

Resumen

El aprendizaje de políticas de robots basado en visión, que mapea entradas visuales a acciones, requiere una comprensión holística de diversas tareas visuales más allá de las necesidades de una sola tarea como clasificación o segmentación. Inspirados en esto, presentamos Theia, un modelo de base visual para el aprendizaje de robots que destila múltiples modelos de base visual listos para usar entrenados en diversas tareas visuales. Las representaciones visuales ricas de Theia codifican conocimientos visuales diversos, mejorando el aprendizaje de robots en etapas posteriores. Experimentos extensos demuestran que Theia supera a sus modelos maestros y a modelos previos de aprendizaje de robots utilizando menos datos de entrenamiento y tamaños de modelo más pequeños. Además, cuantificamos la calidad de las representaciones visuales pre-entrenadas y postulamos que una mayor entropía en las distribuciones de normas de características conduce a un mejor rendimiento en el aprendizaje de robots. El código y los modelos están disponibles en https://github.com/bdaiinstitute/theia.

English

Vision-based robot policy learning, which maps visual inputs to actions, necessitates a holistic understanding of diverse visual tasks beyond single-task needs like classification or segmentation. Inspired by this, we introduce Theia, a vision foundation model for robot learning that distills multiple off-the-shelf vision foundation models trained on varied vision tasks. Theia's rich visual representations encode diverse visual knowledge, enhancing downstream robot learning. Extensive experiments demonstrate that Theia outperforms its teacher models and prior robot learning models using less training data and smaller model sizes. Additionally, we quantify the quality of pre-trained visual representations and hypothesize that higher entropy in feature norm distributions leads to improved robot learning performance. Code and models are available at https://github.com/bdaiinstitute/theia.