PolyTask: Learning Unified Policies through Behavior Distillation

TL;DR

PolyTask addresses the challenge of learning a single unified policy capable of solving many embodied tasks by first training task-specific policies with demonstrations and RL, then distilling these into a single conditioned policy via Behavior Distillation. The method enables lifelong learning through offline distillation that does not require concurrent task environments, and it preserves a constant parameter count while mitigating forgetting. Empirical results across three simulated suites and a real-robot setup show PolyTask outperforms prior multi-task and lifelong-learning baselines, with strong robustness to design choices such as rewards, network size, and goal conditioning modalities. The work highlights practical pathways for scalable, lifelong embodied AI, while noting storage and forward-transfer aspects as avenues for future improvement.

Abstract

Unified models capable of solving a wide variety of tasks have gained traction in vision and NLP due to their ability to share regularities and structures across tasks, which improves individual task performance and reduces computational footprint. However, the impact of such models remains limited in embodied learning problems, which present unique challenges due to interactivity, sample inefficiency, and sequential task presentation. In this work, we present PolyTask, a novel method for learning a single unified model that can solve various embodied tasks through a 'learn then distill' mechanism. In the 'learn' step, PolyTask leverages a few demonstrations for each task to train task-specific policies. Then, in the 'distill' step, task-specific policies are distilled into a single policy using a new distillation method called Behavior Distillation. Given a unified policy, individual task behavior can be extracted through conditioning variables. PolyTask is designed to be conceptually simple while being able to leverage well-established algorithms in RL to enable interactivity, a handful of expert demonstrations to allow for sample efficiency, and preventing interactive access to tasks during distillation to enable lifelong learning. Experiments across three simulated environment suites and a real-robot suite show that PolyTask outperforms prior state-of-the-art approaches in multi-task and lifelong learning settings by significant margins.

Figures (13)

• Figure 1: PolyTask is a technique to train a single unified policy $\Pi$ that can solve a range of different tasks. This is done by first learning a single-task policy for each task followed by distilling it into a unified policy. Distillation allows our unified policy to assimilate additional tasks in a lifelong-learning fashion without needing to increase the parameter count of the policy. Once trained, the unified policy can solve tasks by conditioning on task identifiers such as goal image, text description, or one-hot labels.
• Figure 2: An illustration of the difference between training the unified policy on multi-task and lifelong learning settings.
• Figure 3: PolyTask is evaluated across 3 simulated benchmarks - the DeepMind Control suite, the Meta-World benchmark, and the Franka kitchen environment.
• Figure 4: A comparison between the performance of Fine-tuning julian2020efficientxie2022lifelong and PolyTask on the Meta-World benchmark [(a), (b)] and the Franka kitchen environment [(c), (d)] in a lifelong learning setting. We observe that PolyTask exhibits a significantly better ability to tackle catastrophic forgetting.
• Figure 5: Pixel-based evaluation for lifelong learning on 16 tasks in Meta-World, and 6 tasks in Franka kitchen.