Source-linked AI summary
Multi-Task Reinforcement Learning with Soft Modularization
Ruihan Yang, Huazhe Xu, Yi Wu, Xiaolong Wang
TL;DR
Multi-task RL must share parameters across tasks while avoiding unclear reuse and interfering gradients. The paper uses a task-conditioned routing network to softly combine modules in a base policy, and reports better efficiency and performance across robotic manipulation benchmarks, especially as task diversity grows.
Problem
Multi-task RL lacks a clear way to reuse parameters across tasks while managing interference among jointly trained task gradients.
Method
Soft modularization combines shared base-network modules using task- and state-conditioned routing weights rather than explicitly specifying policy structure.
Results
The method improves sample efficiency and success rate over baselines, with advantages becoming more pronounced for diverse 50-task settings.
Takeaways & Limitations
Soft modularization supports effective sharing and reuse of network components across tasks and may enable future zero-shot generalization to unseen tasks.
Takeaways & Limitations
Hierarchical alternatives face joint-training challenges, while separate sub-policy training can require predefined subtasks or sophisticated subgoal discovery that is typically infeasible in real-world applications.
Abstract
from arXiv · showhide
Multi-task learning is a very challenging problem in reinforcement learning. While training multiple tasks jointly allow the policies to share parameters across different tasks, the optimization problem becomes non-trivial: It remains unclear what parameters in the network should be reused across tasks, and how the gradients from different tasks may interfere with each other. Thus, instead of naively sharing parameters across tasks, we introduce an explicit modularization technique on policy representation to alleviate this optimization issue. Given a base policy network, we design a routing network which estimates different routing strategies to reconfigure the base network for each task. Instead of directly selecting routes for each task, our task-specific policy uses a method called soft modularization to softly combine all the possible routes, which makes it suitable for sequential tasks. We experiment with various robotics manipulation tasks in simulation and show our method improves both sample efficiency and performance over strong baselines by a large margin.
1 Introduction
Multi-task RL aims to share skills across diverse robotic tasks, but joint optimization can cause task interference. The paper introduces soft modularization, which routes task-conditioned policies through weighted combinations of shared modules and improves efficiency and performance.
- Multi-task RL can improve sample efficiency through shared components and provide a curriculum from easier to harder tasks.
- Jointly training many tasks is difficult because task interactions are unclear and optimization for one task can harm others.
- Hierarchical modular approaches require jointly optimizing sub-policies and high-level policies, while separate sub-policy training may require predefined subtasks or discovered subgoals.
- Soft modularization automatically combines modules for each task using a base policy network and a routing network conditioned on task embeddings and state.
- The routing network outputs assignment probabilities, allowing task-specific base networks to use weighted shared modules and enabling joint backpropagation through routing weights.
- In Meta-World's 50 robotic manipulation tasks, the method improves sample efficiency and final performance, nearly doubling success rate over multi-task baselines in the 50-task setting.
2 Related Work
Prior multi-task methods address task interference through gradient manipulation or compositional architectures. Soft or hard routing and mixture-of-experts approaches motivate modular sharing, but hard routing is unstable for sequential RL.
- Gradient-based methods balance or project task gradients, but gradient-similarity optimization can be unstable under within-task gradient variance.
- Compositional models reduce gradient interference by allowing different tasks to use different modules, supporting broader generalization.
- Hard routing policies are difficult for sequential RL because jointly training control and routing policies produces high-variance policy gradients and training instability.
- Mixture-of-experts methods commonly select among experts, whereas this work uses multilayer modules with multiple routing decisions guided by a routing network.
3 Background
The paper formulates each task as a finite-horizon continuous MDP and uses SAC to learn policies balancing task performance with entropy. It extends SAC to multi-task learning with a task-conditioned policy optimized across sampled tasks.
- Each task is modeled as a finite-horizon MDP with continuous states and actions, stochastic transitions, rewards, horizon H, and discount factor γ.
- SAC is an off-policy actor-critic method whose actor seeks task success while maintaining randomness, with policy, Q-function, and temperature parameters optimized.
- The temperature α acts as an entropy penalty coefficient and is learned to maintain a desired minimum expected entropy.
- Multi-task SAC learns a single task-conditioned policy π(a|s, z), where z is a task embedding, and maximizes average expected return over tasks sampled from p(T).
- The multi-task Q-function objective averages task-specific objectives JQ,T(θ) over the task distribution.
4 Method
The method uses a modular base policy and a routing network to learn task-specific soft combinations of shared modules. It also introduces task-dependent objective weighting to balance learning speeds across tasks.
- Soft Modularization: Soft modularization combines shared modules probabilistically instead of selecting discrete routes for each task.The differentiable routing weights permit direct backpropagation and joint training of the base and routing networks.
- Network Architecture: The policy receives the current state and task embedding, producing state and task representations used by both the modules and routing network.The task embedding is represented with a one-hot task vector followed by a fully connected layer.
- Routing Network: The routing network generates n × n module-connection probabilities across L − 1 routing layers for a base network with L module layers and n modules per layer.Its probabilities are computed from prior routing information, the state representation, and the task representation, then normalized with softmax.
- Base Policy Network: The base policy contains L layers of n modules, whose outputs are weighted and combined according to the routing probabilities to produce the action distribution.The Q-function uses a similar soft-modularized architecture, but its base and routing weights are trained independently from the policy network.
- Task-Balanced Training: Task-specific objective weights are learned to balance training, becoming smaller for high policy confidence and larger for low confidence.The weights adjust both policy and Q-function objectives, with task-specific temperature parameters linked to the SAC entropy mechanism.
5 Experiments
Experiments in Meta-World evaluate soft modularization across 10- and 50-task robotics benchmarks, comparing it with multi-task, mixture-of-experts, hard-routing, and single-task baselines. The method improves convergence, sample efficiency, and success rates, while routing visualizations show task differentiation and skill reuse.
- Experimental setup: Meta-World provides MT10 and MT50 benchmarks with fixed-goal and goal-conditioned robotic manipulation tasks, evaluated using average cross-task success rate.Experiments use SAC, three random seeds, and training curves reporting variance; training budgets differ across baselines and methods.
- Routing Network Visualization: Routing visualizations show similar module connections across some tasks and distinct routing-probability clusters across different tasks.These patterns indicate shared skills can be reused while task-specific skillsets remain distinguishable.
- Quantitative Results: 2% better final success rate than the best baseline on MT10-Fixed, with faster convergence despite the modest final-performance difference.The paper attributes the limited final gain to the relative simplicity of training 10 fixed-goal tasks.
- Quantitative Results: Around 24% improvement over baselines on both MT50-Fixed and MT50-Conditioned, with Ours (Deep) outperforming Ours (Shallow) in MT50.The authors suggest more complex routing provides additional choices and reduces harmful task interactions as task count increases.
- Effects on Network Capacity: The smallest model is around 10% better than MT-MH-SAC-5-Wide despite that baseline using 4.2x more parameters, while larger baseline capacity quickly saturates.Training curves also show faster convergence and better performance for the proposed method.
- Comparison with Single Task Policy: With 15 million samples on MT10-Conditioned, Ours (Shallow) reaches 71.8% average success rate versus 78.5% for average single-task policies.The method approaches single-task performance using fewer examples and fewer parameters through shared skills and routing.
- Analysing Learning Components: Removing balanced training or observation inputs for routing substantially reduces success rate, indicating both components are important to the approach.The observation representation may also help explain the Hard Routing baseline’s poor performance with tabular routing control.
6 Conclusion
The paper proposes multi-task reinforcement learning with soft modularization for robotics manipulation, improving sample efficiency and success rate over baselines, especially as task diversity increases.
- Soft modularization improves sample efficiency and success rate over multi-task reinforcement learning baselines by a large margin.
7 Potential Broader Impact
The work provides a framework for skill and component reuse in multi-task reinforcement learning, with potential applications to zero-shot skill transfer and more generalizable robots.
- The learned skill modules provide a framework for skill and component reuse in multi-task reinforcement learning.
- The framework may inspire zero-shot skill transferring and sharing across tasks.
- Improved sample efficiency could make reinforcement learning more accessible and reduce the environmental cost of training.