Research map / Reinforcement learning / Imitation learning
Active Fine-Tuning of Multi-Task Policies
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes AMF, an active task-selection method that picks which tasks to demonstrate during multi-task policy fine-tuning by maximizing information gain about the expert, with convergence guarantees and an adaptive prior against forgetting.
Problem
Given a pre-trained multi-task generalist policy and a limited budget of expert demonstrations, which tasks should be demonstrated (and how often) to maximize multi-task performance after behavioral-cloning fine-tuning.
Approach
Select each demonstration task by maximizing expected information gain about the expert policy over visited states (approximated via importance-weighted occupancy estimates), with a GP / loss-gradient-embedding entropy estimate for neural policies, plus an Adaptive Prior that mixes actions from a frozen pre-trained copy to reduce catastrophic forgetting.
Key results
- On the 2D integrator with mismatched pre-training, AMF-GP's advantage over uniform sampling grows as the pre-training distribution covers fewer tasks (e.g., 6/12 tasks).
- AMF-NN matches uniform sampling when pre-training is uniform across tasks and outperforms it under mismatched pre-training on Metaworld and FrankaKitchen, from state and RGB inputs (means with 90% bootstrap CIs over 10 seeds).
- Last-layer embedding GP and test-time dropout entropy selection yield lower multi-task performance than loss-gradient-embedding AMF-NN.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use AMF-GP instead of uniform task sampling | GP policy fine-tuning · 2D integrator, mismatched pre-training distribution | multi-task return | Helped |
| Use AMF-GP instead of uniform task sampling | GP policy fine-tuning · 2D integrator, matched pre-training distribution | multi-task return | Helped |
| Use AMF-NN instead of uniform task selection | Neural multi-task policies · Meta-World and FrankaKitchen, mismatched pre-training | success rate | Helped |
| Use AMF-NN instead of uniform task selection | Neural multi-task policies · Meta-World and FrankaKitchen, matched pre-training | success rate | No clear effect |
| Add Adaptive Prior during fine-tuning | Neural multi-task policies · Meta-World and FrankaKitchen | success rate | Helped |
| Add Adaptive Prior during fine-tuning | Fine-tuned policies · Robomimic | success rate | Helped |
| Replace loss-gradient embeddings with last-layer embeddings for entropy estimation (ablation) | AMF-NN · Neural multi-task policy experiments | multi-task performance | Hurt |
| Use test-time Dropout for entropy estimation instead of loss-gradient embeddings (ablation) | AMF-NN · Neural multi-task policy experiments | multi-task performance | Hurt |
Surprising findings
- Uniform task sampling without an adaptive prior suffers performance degradation from catastrophic forgetting of pre-training demonstrations.
- Regularization-based forgetting mitigation was not found to be empirically effective.
Limitations
- Performance of neural uncertainty estimation (loss-gradient embeddings) is an open problem; relies on approximations.
- Gains are modest when the pre-trained policy is equally capable across tasks and uniform sampling suffices.
- No large-scale generalist policy evaluation due to lack of open calibrated benchmarks.
Experimental setup
Datasets: Metaworld FrankaKitchen Robomimic
Benchmarks: 2D integrator Metaworld FrankaKitchen Robomimic
Metrics: multi-task success rate returns performance gap to expert
Baselines: uniform random task sampling last-layer embedding GP test-time dropout entropy
Builds on: behavioral cloning active learning via information gain (GP-UCB style analysis) loss gradient embeddings behavioral priors
Closest papers by meaning
- Task-Agnostic Pre-training and Task-Guided Fine-tuning for Versatile Diffusion Planner (ICML 2025 · similarity 0.88)
- BAKU: An Efficient Transformer for Multi-Task Policy Learning (NeurIPS 2024 · similarity 0.87)
- Self-Improving Embodied Foundation Models (NeurIPS 2025 · similarity 0.87)
- Rapidly Adapting Policies to the Real-World via Simulation-Guided Fine-Tuning (ICLR 2025 · similarity 0.87)
- Diffusion Policy Policy Optimization (ICLR 2025 · similarity 0.87)
- GO-DICE: Goal-Conditioned Option-Aware Offline Imitation Learning via Stationary Distribution Correction Estimation (AAAI 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.