MENTOR: Mixture-of-Experts Network with Task-Oriented Perturbation for Visual Reinforcement Learning

In challenging robotic learning tasks, RL agents are often assigned $K\geq 2$ different tasks or subgoals, each associated with a loss function $L_{i}(\theta)$. The goal is to find optimal agent weights $\theta\in\mathbb{R}^{m}$ that minimize losses across all objectives. In practice, a common approach is to reduce the average loss over all tasks: $\theta^{*}=\arg\min_{\theta\in\mathbb{R}^{m}}\left\{L_{0}(\theta)\overset{\Delta}{=}\frac{1}{K}\sum_{i=1}^{K}L_{i}(\theta)\right\}$. If the agent uses a shared set of parameters $\theta$ (e.g., MLP), meaning all parameters must be simultaneously active to function, the optimization process using gradient descent may compromise individual loss optimization. This issue, known as conflicting gradients (Yu et al., 2020a; Liu et al., 2021), hinders the agent’s ability to optimize its behavior when facing complex scenarios effectively.

To address this issue, we propose to utilize an MoE architecture as a substitute for the MLP backbone in RL agents. The MoE structure is characterized by its composition of modular experts, $\theta_{\mathrm{MoE}}=\{\theta_{1},\theta_{2},\dots,\theta_{N}\}$, which allows the agent to activate different experts via a dynamic routing mechanism flexibly. This enables gradients from different tasks or subgoals to correspond to different sets of parameters. Specifically, the parameters of a given expert are updated only by gradients from similar task scenarios, thereby effectively alleviating the gradient conflict problem.

As illustrated in Figure 2, the MoE agent first processes visual inputs using a CNN-based encoder, transforming them into a latent space $Z$. The router $h$ computes a probability distribution $h(i\mid z)$ over the experts $i$ for a given latent vector $z\in Z$. The top-$k$ experts are selected based on this distribution, and their softmax weights $w_{i}$ are computed. The outputs $a_{i}$ from these top-$k$ experts are weighted combined to produce the final output $a$, as shown in Equations 1 and 2. This MoE structure enables the agent to route input visual features to specialized experts based on specific objectives, optimizing its performance in challenging scenarios such as multi-tasking or multi-stage processes.

To better illustrate the important role of dynamic modular expert learning for RL agents, we conduct a multi-task experiment (MT5) in Meta-World (MW) (Yu et al., 2020b), where the agent ($\texttt{\#}\text{Experts}=16,k=4$) is trained to acquire five opposing skills: Open tasks (Door-Open, Drawer-Open, Window-Open) and Close tasks (Drawer-Close, Window-Close). As shown in Figure 3(a), in addition to sharing some experts for handling common knowledge, the Open and Close tasks have their own dedicated experts (Experts 3, 7 for Open and Experts 9, 10 for Close). We evaluate the cosine similarities (Yu et al., 2020a) for both MLP and MoE agents, as shown in Figure 3(b). The MLP’s gradients show significant conflicts between opposing tasks, resulting in a performance gap (100% success for Close tasks, 82% for Open tasks). In contrast, the MoE model demonstrates higher gradient compatibility, achieving 100% success in both task types.

This structural advantage can also be propagated to challenging single tasks, as the dynamic routing mechanism automatically activates different experts to adjust the agent’s behavior throughout the task, alleviating the burden on shared parameters. We illustrate this through training a same-structure MoE agent on a single, highly challenging Assembly task from MW. Figure 4 shows the engagement of the $k=4$ most active experts during task execution, with Expert 15 serving as the shared module throughout the entire policy execution. The other experts vary and automatically divide the task into four distinct stages: Expert 9 handles gripper control for grasping and releasing; Expert 13 manages arm movement while maneuvering the ring; and Expert 14 oversees the assembly process as the ring approaches its fitting location.

Neural network perturbation is employed to enhance the exploration capabilities in RL. Two key factors influence the effectiveness of this process $\theta_{k}=\alpha\theta_{k-1}+(1-\alpha)\phi,\phi\sim\Phi$. $\alpha$ is the perturbation factor controlling the mix between current agent and perturbation candidate weights. $\phi$ represents the perturbation candidate sampled from a distribution $\Phi$, which typically is a fixed Gaussian noise $\mathcal{N}(\mu,\sigma)$. Previous works (Sokar et al., 2023; Xu et al., 2023; Ji et al., 2024) have investigated the use of the dormant ratio to determine $\alpha$, resulting in improved exploration efficiency (see Section 2). However, the selection of perturbation candidates has not been thoroughly examined. In this work, we propose sampling $\phi$ from a heuristically updated distribution $\Phi_{\text{oriented}}$, generated from past high-performing agents, to provide more task-oriented candidates that better facilitate optimization.

We define $\Phi_{\text{oriented}}$ as a distribution from which the weights of high-performing agents can be sampled. To obtain this distribution, we dynamically maintain a fixed-size set $S_{\text{top}}=\{(\theta,R)\}$, where $(\theta,R)$ denotes an agent with the weights $\theta$ and achieves episode reward $R$. The desired distribution is approximated using a Gaussian distribution, specifically $\Phi_{\text{oriented}}=\mathcal{N}(\mu_{\theta}^{\text{top}},\sigma_{\theta}^{\text{top}})$, where $\mu_{\theta}^{\text{top}}$ and $\sigma_{\theta}^{\text{top}}$ represent the mean and standard deviation of the weights in $S_{\text{top}}$. As shown in Figure 2, the set $S_{\text{top}}$ is updated during training: at episode $t$, if the agent with weights $\theta_{t}$ achieves reward $R_{t}$ higher than the lowest in $S_{\text{top}}$, the tuple $(\theta_{t},R_{t})$ replaces the one with the lowest reward. This update ensures that $\Phi_{\text{oriented}}$ accurately reflects the current set of high-performing agents, thus providing improved perturbation candidates $\phi$ for future iterations. The pseudocode is shown in Algorithm 1.

For illustration, we conduct experiments on the Hopper Hop task from the DeepMind Control Suite (DMC), comparing task-oriented perturbation approach to leading model-free visual RL baselines (DrM (Xu et al., 2023) and DrQ-v2 (Yarats et al., 2021)). Our approach solely replaces DrM’s perturbation mechanism with task-oriented perturbations. Both our method and DrM outperform DrQ-v2 due to dormant-ratio-based perturbation, but our method achieves faster skill acquisition and maintains a lower, smoother dormant ratio throughout training (Figure 5a and 5b). By directly testing perturbation candidates as agents in the task (Figure 5c), we observe that candidates sampled from $\Phi_{\text{oriented}}$ steadily improve throughout training, sometimes even surpassing the performance of the agent they perturb. This demonstrates that $\Phi_{\text{oriented}}$ progressively captures the optimal weight distribution, rather than simply interpolating from past agents, leading to more targeted optimization. In contrast, perturbation candidates from DrM (initialized with Gaussian noise) consistently yield zero reward, indicating the lack of task-relevant information.

Baselines: We compare MENTOR against four leading model-free visual RL methods: DrM (Xu et al., 2023), ALIX (Cetin et al., 2022), TACO (Zheng et al., 2023), and DrQ-v2 (Yarats et al., 2021). DrM, ALIX, and TACO all use DrQ-v2 as their backbone. DrM periodically perturbs the agent’s weights with random noise based on the proportion of dormant neurons in the neural network; ALIX adds regularization to the encoder gradients to mitigate overfitting; and TACO employs contrastive learning to improve latent state and action representations.

Experimental Settings: We evaluate MENTOR on a diverse set of tasks across three simulation environments with complex dynamics and even sparse reward. The DMC includes challenging tasks like Dog Stand, Dog Walk, Manipulator Bring Ball, and Acrobot Swingup (Sparse), focusing on long-horizon continuous locomotion and manipulation challenges. The MW environment provides a suite of robotic tasks including Assembly, Disassemble, Pick Place, Coffee Push (Sparse), Soccer (Sparse), and Hammer (Sparse), which test the agent’s manipulation abilities and require sequential reasoning. The Adroit environment includes complex robotic manipulation tasks such as Door and Hammer, which involve controlling dexterous hands to interact with articulated objects. Notably, DMC tasks are evaluated using episode reward, while tasks in MW and Adroit are assessed based on success rate. For each method on each task, we conducted experiments using four random seeds; detailed hyperparameters and training settings are provided in Appendix B.

Results: Figure 6 presents performance comparisons between MENTOR and the baselines. In the DMC tasks, Dog Stand and Dog Walk feature high action dimensionality with a 38-dimensional action space representing joint controls for the dog model. These tasks also have complex kinematics involving intricate joint coordination, muscle dynamics, and collision handling, making them challenging to optimize. Our method outperforms the top baseline, achieving approximately 17% and 10% higher episode rewards, respectively. In the MW tasks, the Hammer (Sparse) task stands out. It requires a robotic arm to hammer a nail into a wall, with highly sparse rewards: success yields significantly larger rewards than merely touching or missing the nail. In fact, the reward for failure is only one-thousandth of the success reward, making the task extremely sparse. However, our task-oriented perturbation effectively captures these sparse rewards, reducing the required training frames by 70% compared to the best baseline. In the Adroit tasks, our method achieves nearly 100% success with significantly less training time, while the most competitive counterpart (DrM) requires more frames, and other baselines fail to match performance even after 6 million frames. A key highlight is the Door task, which involves multiple stages of dexterous hand manipulation—grasping, turning, and opening the door. Leveraging the MoE architecture, our method reduces training time to achieve over 80% success by approximately 23% compared to the best baseline. In summary, MENTOR demonstrates superior efficiency and performance compared to the strongest existing model-free visual RL baselines across all 12 tasks,.

Our real-world RL experiments evaluate the practical applicability of MENTOR in robotic manipulation tasks. We design three tasks to highlight key challenges in real-world robotics: multi-task learning, multi-stage deformable object manipulation, and dynamic skill acquisition.

Experimental Settings: All tasks use a Franka Panda arm for execution and RealSense D435 cameras for RGB visual observations, which include both overall and close-up views to capture global and local information. The reward functions are based on the absolute distance between the current and desired states. To prevent trajectory overfitting, the end-effector’s initial position is randomly sampled from a predefined region at the start of each episode. Tasks are described below and shown in Figure 7. Further details can be found in Appendix C.

Peg Insertion: This task simulates an assembly-line scenario where fine-grained insertion of various objects are required. The agent needs to develop multi-task learning skills to insert pegs with three different shapes (Star, Triangle, and Arrow) into corresponding sockets. Training such agents in simulators is difficult due to the complexities of contact-rich interactions, making this task ideal for real-world reinforcement learning and evaluation.

Cable Routing: Manipulating deformable cables presents significant challenges due to the complexities of modeling and simulating their physical dynamics, making this task ideal for direct, model-free visual RL training in real-world environments. In this scenario, the robot must guide a cable into two parallel slots. Since both slots cannot be filled simultaneously, the agent must perform the task sequentially, requiring long-horizon, multi-stage planning to successfully accomplish the task.

Tabletop Golf: In this task, the robot uses a golf club to strike a ball on a grass-like surface, aiming to land it in a target hole. An automated reset system retrieves the ball when it reaches the hole, enters a mock water hazard, or rolls out of bounds, and randomly repositions it. The agent must learn to approach the ball, control the club’s striking force and direction to guide the ball toward the hole while avoiding obstacles through real-world interaction.

Results: Our policies demonstrate robust performance during evaluation as shown in Figure 7. In Peg Insertion, the agent randomly selects a peg from the shelf and inserts it from varying initial positions. It gradually learns to align the peg shape with the corresponding hole and adjust the angle for accurate insertion. During one execution, as the peg nears the hole, we manually disturb by altering the robot arm’s pose significantly. Despite this interference, the agent successfully completes the task relying solely on visual observations. In Cable Routing, where the cable cannot be placed into parallel slots simultaneously, the agent learns to prioritize routing it into the farther slot first, then into the closer one. This second step requires careful handling to avoid dislodging the cable from the first slot. During execution, if the cable is randomly removed from the slot, the agent can visually detect this issue and re-route it back into position. In Tabletop Golf, the agent must master two key skills: striking the ball with the correct direction and force, and repositioning the club to follow the ball after the strike. Due to a "water hazard", the ball cannot be struck directly toward the target hole from its starting position. The agent learns to angle its shots to bypass the hazard and guide the ball into the hole. No interference is applied during this task, as the ball’s rolling on the grass-like surface introduces sufficient variability.

Ablation Study: We conduct a detailed ablation study to demonstrate the effectiveness of MENTOR in improving sample efficiency and performance, as shown in Table 1.

The first two rows reveal that utilizing the pretrained visual encoder (Lin et al., 2024) instead of a CNN trained from scratch results in an average performance improvement of 9%. However, no significant performance gain is observed in simulation benchmarks with this substitution. This discrepancy may arise from the gap between simulation and real-world environments, where real scenes offer richer textures more aligned with the pretraining domain.

Furthermore, the results confirm the effectiveness of our technical contributions. When the MoE structure is removed from the agent (i.e., replaced with an MLP, as in MENTOR w/o MoE), overall performance drops by nearly 30%. Additionally, further switching the task-oriented perturbation mechanism to basic random perturbation (as in DrM) leads to an additional performance decline of approximately 30%. We further extend the training process of the DrM baseline to reach the same performance level as MENTOR, with the training time comparison shown in Figure 8, which demonstrates an average 37% improvement in time efficiency for our method. These findings underscore the importance of each component in achieving superior results.