Multi-Task Reinforcement Learning with Soft Modularization

TL;DR

This work tackles the optimization challenges of multi-task reinforcement learning by introducing soft modularization: a base policy network composed of modules whose connections are softly weighted by a routing network conditioned on the task, enabling end-to-end training without predefined hierarchies. By jointly learning module routing and task-conditioned policies, the method achieves superior sample efficiency and final performance on a large suite of robotic manipulation tasks in Meta-World, particularly as task count increases. The paper also provides extensive analyses of routing visualizations, capacity effects, and ablations that underscore the importance of both balanced multi-task optimization and observation-informed routing. Overall, soft modularization demonstrates strong cross-task transfer and opens avenues for zero-shot generalization to unseen tasks with efficient training and scalable modular designs.

Abstract

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.

Figures (6)

• Figure 1: Our multi-task policy network with soft modularization. Given different tasks, our network generate different soft combination of network modules. Gray squares represent network modules and red lines represent the connection between modules (Darker red indicates larger weight).
• Figure 2: Our framework contains a base policy network with multiple modules (left) and a routing network (right) generating connections between modules in the base policy network.
• Figure 3: Sampled observation and corresponding routing. Each column shows two different tasks sharing similar routing. The shared parts are highlighted with blue boxes.
• Figure 4: Probabilities from the routing network for different tasks are extracted and visualized with t-NSE. Routing probabilities from different tasks are grouped in different clusters.
• Figure 5: Training curves of different methods on all benchmarks (Concrete lines: the average over 3 seeds; Shaded areas: the standard deviation over 3 seeds). For MT10, our method converges much faster than the baselines. For MT50, we achieve a large gain on sample efficiency and performance.