Task Tokens: A Flexible Approach to Adapting Behavior Foundation Models

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Vainshtein, Ron, Rimon, Zohar, Mannor, Shie, Tessler, Chen
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866912984058560512
author Vainshtein, Ron
Rimon, Zohar
Mannor, Shie
Tessler, Chen
author_facet Vainshtein, Ron
Rimon, Zohar
Mannor, Shie
Tessler, Chen
contents Recent advancements in imitation learning have led to transformer-based behavior foundation models (BFMs) that enable multi-modal, human-like control for humanoid agents. While excelling at zero-shot generation of robust behaviors, BFMs often require meticulous prompt engineering for specific tasks, potentially yielding suboptimal results. We introduce "Task Tokens", a method to effectively tailor BFMs to specific tasks while preserving their flexibility. Our approach leverages the transformer architecture of BFMs to learn a new task-specific encoder through reinforcement learning, keeping the original BFM frozen. This allows incorporation of user-defined priors, balancing reward design and prompt engineering. By training a task encoder to map observations to tokens, used as additional BFM inputs, we guide performance improvement while maintaining the model's diverse control characteristics. We demonstrate Task Tokens' efficacy across various tasks, including out-of-distribution scenarios, and show their compatibility with other prompting modalities. Our results suggest that Task Tokens offer a promising approach for adapting BFMs to specific control tasks while retaining their generalization capabilities.
format Preprint
id arxiv_https___arxiv_org_abs_2503_22886
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Task Tokens: A Flexible Approach to Adapting Behavior Foundation Models
Vainshtein, Ron
Rimon, Zohar
Mannor, Shie
Tessler, Chen
Machine Learning
Robotics
Recent advancements in imitation learning have led to transformer-based behavior foundation models (BFMs) that enable multi-modal, human-like control for humanoid agents. While excelling at zero-shot generation of robust behaviors, BFMs often require meticulous prompt engineering for specific tasks, potentially yielding suboptimal results. We introduce "Task Tokens", a method to effectively tailor BFMs to specific tasks while preserving their flexibility. Our approach leverages the transformer architecture of BFMs to learn a new task-specific encoder through reinforcement learning, keeping the original BFM frozen. This allows incorporation of user-defined priors, balancing reward design and prompt engineering. By training a task encoder to map observations to tokens, used as additional BFM inputs, we guide performance improvement while maintaining the model's diverse control characteristics. We demonstrate Task Tokens' efficacy across various tasks, including out-of-distribution scenarios, and show their compatibility with other prompting modalities. Our results suggest that Task Tokens offer a promising approach for adapting BFMs to specific control tasks while retaining their generalization capabilities.
title Task Tokens: A Flexible Approach to Adapting Behavior Foundation Models
topic Machine Learning
Robotics
url https://arxiv.org/abs/2503.22886