# Task Tokens: A Flexible Approach to Adapting Behavior Foundation Models

arXiv: https://arxiv.org/abs/2503.22886 · HTML: https://mlsanity.com/paper/2503.22886/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Robot manipulation imitation learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 18 of its 29 references that are in the map.
- First posted on arXiv 28 Mar 2025.

## Placed in

- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 18 of 18)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Parameter-Efficient Transfer Learning for NLP](https://mlsanity.com/paper/1902.00751/) (2019)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [An Empirical Investigation of Catastrophic Forgetting in Gradient-Based Neural Networks](https://mlsanity.com/paper/1312.6211/) (2013)
- [Perpetual Humanoid Control for Real-time Simulated Avatars](https://mlsanity.com/paper/2305.06456/) (2023)
- [Universal Humanoid Motion Representations for Physics-Based Control](https://mlsanity.com/paper/2310.04582/) (2023)
- [MaskedMimic: Unified Physics-Based Character Control Through Masked Motion Inpainting](https://mlsanity.com/paper/2409.14393/) (2024)
- [Learning One Representation to Optimize All Rewards](https://mlsanity.com/paper/2103.07945/) (2021)
- [CALM: Conditional Adversarial Latent Models for Directable Virtual Characters](https://mlsanity.com/paper/2305.02195/) (2023)
- [ExBody2: Advanced Expressive Humanoid Whole-Body Control](https://mlsanity.com/paper/2412.13196/) (2024)
- [Mobile-TeleVision: Predictive Motion Priors for Humanoid Whole-Body Control](https://mlsanity.com/paper/2412.07773/) (2024)
- [CooHOI: Learning Cooperative Human-Object Interaction with Manipulated Object Dynamics](https://mlsanity.com/paper/2406.14558/) (2024)
- [Masked Trajectory Models for Prediction, Representation, and Control](https://mlsanity.com/paper/2305.02968/) (2023)
- [SMPLOlympics: Sports Environments for Physically Simulated Humanoids](https://mlsanity.com/paper/2407.00187/) (2024)

## Built on by (most cited first, 3 of 3)

- [InterPrior: Scaling Generative Control for Physics-Based Human-Object Interactions](https://mlsanity.com/paper/2602.06035/) (2026)
- [Zero-Shot Off-Policy Learning](https://mlsanity.com/paper/2602.01962/) (2026)
- [Progressively Learning Heterogeneous Skills in a Unified Latent Space](https://mlsanity.com/paper/2608.23258/) (2026)

## Abstract (arXiv)

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.