Source-linked AI summary
T3S: Improving Multi-Task Reinforcement Learning with Task-Specific Feature Selector and Scheduler
Yuanqiang Yu, Tianpei Yang, Yongliang Lv, Yan Zheng, Jianye Hao
TL;DR
MTRL must share knowledge across tasks without allowing destructive interference from inappropriate parameter sharing. T3S uses hypernetwork-generated soft masks for fine-grained feature sharing and schedules tasks by progress and learning speed; experiments report significant gains across robotics manipulation tasks.
Problem
MTRL shares representations to improve efficiency, but destructive interference remains because task interactions and appropriate parameter sharing are not controlled.
Method
T3S combines hypernetwork-based task-specific feature selectors with a scheduler whose task-selection probabilities are inversely proportional to task progress and learning speed.
Results
T3S significantly outperforms state-of-the-art MTRL algorithms on various robotics manipulation tasks.
Takeaways & Limitations
Fine-grained feature sharing and difficulty-aware task scheduling promote knowledge transfer between similar tasks while addressing inter-task interference.
Abstract
from arXiv · showhide
Multi-task reinforcement learning (MTRL) is a technique to train multiple tasks simultaneously, where previous works usually train a single model to solve different tasks by sharing parameters across various tasks. However, these methods are faced with inter-task interference since what parameters should be shared across tasks is not addressed, dramatically reducing learning efficiency. To solve these problems, we propose a novel MTRL framework called Task-Specific feature Selector and Scheduler (T3S), which consists of two components: a feature selector and a task scheduler. Specifically, the feature selectors employ hypernetworks to construct task-specific soft masks, which can be applied by globally shared representation to construct task-specific features. The task scheduler selects tasks for learning through two metrics, where the selection probability is inversely proportional to task progress (e.g., success rate) and task learning speed. Experimental results show that T3S consistently outperforms the state-of-the-art MTRL algorithms on various robotics manipulation tasks.
I. INTRODUCTION
MTRL improves efficiency by sharing representations across tasks, but destructive interference arises when task interactions and parameter sharing are not controlled. T3S addresses this with fine-grained feature selection and scheduling of challenging tasks.
- Motivation: MTRL shares representations across tasks to improve learning efficiency, but different tasks can interfere during joint training.Single-task learning requires separate agents and prevents knowledge transfer during training.
- Motivation: Destructive interference occurs because shared-network updates can improve one task while impairing another.The issue reflects uncertainty about how tasks affect one another in a single model.
- T3S: T3S combines a feature selector and task scheduler to promote fine-grained knowledge sharing between tasks.The feature selector addresses parameter sharing at the feature level, while the scheduler addresses task-selection imbalance.
- T3S: The feature selector uses hypernetworks and task IDs to construct task-specific soft masks over globally shared representations.This enables selective inter-task parameter sharing rather than uniformly sharing all features.
- T3S: Task-selection probability is inversely proportional to task progress and task learning speed, prioritizing less-progressed and slower-learning tasks.This scheduler is intended to prevent easy tasks from receiving inappropriate training focus.
- Results: T3S significantly outperforms state-of-the-art MTRL algorithms on various robotics manipulation tasks and can combine with off-policy DRL algorithms.The reported evaluation covers robotics manipulation settings.
II. BACKGROUND
The paper formulates each task as a finite MDP and uses SAC as the off-policy reinforcement-learning foundation. Tasks may have different MDPs and separate entropy temperatures.
- Problem Settings: Each reinforcement-learning task is modeled as a finite MDP defined by state and action spaces, transitions, rewards, horizon, and discount factor.The objective is to learn a policy maximizing expected discounted return.
- MTRL Setting: The MTRL setting treats a collection of tasks equally, although each task may have a different MDP.Meta-World provides examples such as opening a door and closing a window.
- Soft Actor-Critic: SAC trains an actor to balance expected return against action-distribution entropy.The temperature parameter α controls policy stochasticity and maintains the entropy level.
- Soft Actor-Critic: SAC updates the value network using a Bellman target and temporal-difference error.The target value uses a sampled next action and a target network.
- Task-Specific SAC: T3S assigns a separate temperature α to each task while retaining the standard SAC policy and critic optimization approach.This accommodates different task learning statuses.
III. METHODOLOGY
T3S combines a soft-mask network architecture with a task scheduler, then integrates the framework with SAC to improve sample efficiency and reduce inter-task interference.
- Network Architecture: T3S uses feature selectors and hypernetworks to learn task-specific soft masks over a globally shared network.This enables inter-task parameter sharing at the feature level.
- Task Scheduler: The task scheduler focuses training on complex tasks through a scheduling mechanism based on task metrics.The stated goals are improved sample efficiency and reduced inter-task interference.
- SAC Integration: T3S is combined with SAC as the specific DRL algorithm used in the framework description.The methodology presents the architecture, scheduler, and SAC integration as separate components.
A. Task-Specific Feature Selector
The task-specific feature selector derives soft masks from task context and shared features, then applies them to globally shared representations to produce task-specific features.
- Architecture: The architecture contains a globally shared network and several feature selectors that extract task-specific features using soft masks.The number of feature selectors can be defined according to task complexity.
- Feature Filtering: Feature selectors apply learned soft masks to globally shared features, filtering features in a fine-grained manner.The resulting task-specific features are passed to the next feature selector.
- Mask Generation: Hypernetworks take task context, such as a one-hot task ID, and generate mask-generator weights.The mask generator then takes shared features as input and outputs the task-specific soft mask.
- Mask Semantics: The sigmoid activation controls the degree of feature sharing in each soft mask.Values tending toward 1 indicate shared features, whereas values tending toward 0 indicate task-specific features.
- End-to-End Learning: Shared-network features and soft masks are trained jointly to optimize task-specific performance while mitigating inter-task interference.This couples representation learning and task-specific filtering in an end-to-end process.
B. Task Scheduler
T3S schedules task sampling using both relative task progress and learning speed, prioritizing tasks that are progressing slowly. It updates sampling probabilities during evaluation and trains on sampled task subsets.
- Motivation: The scheduler is designed to avoid inefficiently sampling simple tasks after they converge.Uniform sampling can continue allocating training to already-converged tasks, whereas optimized scheduling focuses on task progress and learning speed.
- Scheduler procedure: At each evaluation interval, T3S evaluates every task and updates the task sampling distribution.During training, it samples a subset of K tasks according to the updated distribution, collects trajectories, and stores them in the replay buffer.
- Scheduler procedure: The training loop updates actor and critic parameters with LSAC and then updates the target networks from minibatches sampled from the replay buffer.This procedure follows task evaluation and subset sampling.
- Task scheduling principle: T3S assigns higher sampling probability to tasks with slower progress and learning speed.Task progress can use performance such as normalized total reward or success rate; learning speed uses performance increments over a time interval.
- Task scheduling principle: The scheduler combines task metrics through weighted averaging before computing the sampling distribution.The temperature τ shapes the distribution, while αk weights metric c^(k).
C. T3S-SAC
T3S-SAC integrates T3S with SAC by updating task probabilities during evaluation, sampling task subsets for parallel data collection, and training from a shared replay buffer. Replay samples containing all tasks help avoid catastrophic forgetting of well-learned tasks.
- T3S-SAC procedure: T3S-SAC initializes a task distribution, updates it from task metrics during evaluation, and samples task subsets during training.The sampled tasks collect trajectories in parallel before reinforcement-learning updates are performed.
- T3S-SAC procedure: T3S-SAC stores trajectories from sampled tasks in a replay buffer and updates the agent and target networks using RL loss.The framework is combined with SAC as the underlying deep reinforcement-learning algorithm.
- T3S-SAC procedure: Replay-buffer samples usually contain all tasks, which can avoid catastrophic forgetting for well-learned tasks.This allows training to remain focused on selected tasks while retaining experience from other tasks.
A. Experimental Results
T3S is evaluated on Meta-World manipulation benchmarks against representative multi-task reinforcement-learning baselines. Across MT10 and MT50 settings, it improves sample efficiency and performance, while its fine-grained sharing and scheduler support lower-cost training than CAGrad.
- Evaluation setup: T3S is evaluated on Meta-World’s MT10 and MT50 robotics manipulation benchmarks with fixed and randomly initialized goals.The benchmarks contain 10 or 50 tasks, and all algorithms are trained from scratch.
- Evaluation setup: T3S is compared with MT-SAC, MT-MH-SAC, MMoE, Soft Module, and CAGrad using Adam with a 3×10−4 learning rate.MT10 experiments run for 15 million steps with batch size 1280.
- MT10 and MT50 results: T3S significantly improves sample efficiency and reduces inter-task interference relative to all baselines on MT10, including randomly generated-goal settings.Success rates are averaged over the last 200,000 training steps across 3 seeds.
- MT10 and MT50 results: T3S greatly outperforms other baselines on MT50 except CAGrad, matching CAGrad’s final success rate while showing better sample efficiency and lower variance.The reported advantages are attributed to fine-grained sharing and task scheduling.
- Mechanistic analysis: The scheduler concentrates sampling on difficult tasks, while t-SNE visualizations place tasks with similar skills closer together through their soft masks.In MT10-RAND, push and pick-place receive sustained focus, while mastered tasks are sampled less often.
- Computational cost: T3S is around 20x faster than CAGrad per update step because CAGrad requires many extra optimization steps.The comparison focuses on CAGrad because other baselines have no significant training-time difference from T3S.
B. Ablation Studies
Ablations examine the scheduler and feature selector separately and in combination with existing baselines. Removing or simplifying these components degrades performance, while adding the scheduler improves baseline sampling efficiency.
- Evaluation protocol: The ablations report average success rates on MT10-RAND and MT10-FIXED over the last 200,000 steps across 3 seeds.Figure 7 presents scheduler and feature-selector comparisons, while Tables III and IV provide quantitative results.
- Task scheduler: Removing the task scheduler degrades T3S performance on MT10-RAND, demonstrating that scheduling contributes materially to the framework.The ablation compares the full method with scheduler removal and alternatives excluding task progress or learning speed.
- Task scheduler: The task progress metric produces larger performance gains than the task learning speed metric, suggesting greater weight should be placed on progress.This comparison is made between T3S without progress and T3S without learning speed.
- Task scheduler: Adding the task scheduler improves baseline performance and sampling efficiency compared with uniform sampling.The scheduler addresses imbalanced learning by focusing sampling on more difficult tasks.
- Feature selector: T3S outperforms MT-MH-SAC with the scheduler, supporting the contribution of fine-grained sharing beyond task scheduling alone.Without the feature selector, T3S is equivalent to MT-MH-SAC with the scheduler.
V. RELATED WORK
Related MTRL work addresses task sharing through architecture design or optimization strategies. These approaches include routing, attention, modularization, distillation, and gradient-conflict handling, each with stated trade-offs or limitations.
- Multi-task architectures: Multi-task architecture methods balance shared and task-specific parameters using mechanisms such as MMoE, routing networks, and MTAN.MMoE shares experts with task-specific gating, routing networks dynamically combine function blocks, and MTAN uses a shared network with task-specific attention.
- Multi-task architectures: MMoE weights all expert outputs through gating without distinguishing task-shared from task-specific representations, which can lead to destructive interference.Routing networks dynamically combine function blocks but use reinforcement learning to train the routing policy, which may introduce training instability in temporal RL.
- Optimization strategies: Conflicting gradients provide an optimization view of interference, where following one task’s gradient can reduce another task’s performance when directions oppose.PCGrad uses gradient projection, while CAGrad exploits the worst local improvement of tasks.
- Knowledge distillation: Actor-Mimic and Policy Distillation train task-specific teacher policies before distilling their outputs into a single student policy.These methods use knowledge distillation for multi-task reinforcement learning.
VI. CONCLUSION AND FUTURE WORK
T3S combines fine-grained feature selection with progress- and speed-based task scheduling to improve knowledge transfer in multi-task reinforcement learning. Experiments on robotics manipulation tasks outperform several state-of-the-art MTRL baselines, while evaluation remains simulation-bound.
- Conclusion: T3S uses hypernetworks to generate task-ID-conditioned soft masks that select shared and task-specific features finely.The feature selector automatically determines which features to share and which remain task-specific.
- Conclusion: Its task scheduler selects learning tasks using probabilities inversely proportional to task progress and task learning speed.These two metrics address differences in task difficulty during training.
- Experimental findings: T3S outperforms several state-of-the-art MTRL baselines across various robotics manipulation tasks.The paper also uses t-SNE visualizations to examine the task-distinction capability of the soft masks.
- Future work: The evaluation does not include real robots, and the scheduler’s metric weights are set manually.Future work proposes improving the scheduler and extending T3S to real robots.