Scalable Multi-Objective and Meta Reinforcement Learning via Gradient Estimation

TL;DR

The paper tackles scalable multi-objective reinforcement learning with many tasks by introducing PolicyGradEx, a two-stage method that first learns a meta-policy across all tasks and then uses a first-order surrogate on projected gradients to estimate adaptation to arbitrary subsets. This yields a task-affinity matrix which is clustered into k groups via a convex SDP relaxation, enabling efficient group-wise training. Empirical results on Meta-World and robotics benchmarks show substantial gains over baselines (e.g., ~16-22% average improvements) and up to 26× speedups over full-training approaches, along with Hessian-based insights linking sharpness to generalization in multi-task RL. The work also provides a theoretical Hessian-based generalization bound via PAC-Bayes with anisotropic perturbations, bridging optimization dynamics and generalization in multi-task policy learning.

Abstract

We study the problem of efficiently estimating policies that simultaneously optimize multiple objectives in reinforcement learning (RL). Given $n$ objectives (or tasks), we seek the optimal partition of these objectives into $k \ll n$ groups, where each group comprises related objectives that can be trained together. This problem arises in applications such as robotics, control, and preference optimization in language models, where learning a single policy for all $n$ objectives is suboptimal as $n$ grows. We introduce a two-stage procedure -- meta-training followed by fine-tuning -- to address this problem. We first learn a meta-policy for all objectives using multitask learning. Then, we adapt the meta-policy to multiple randomly sampled subsets of objectives. The adaptation step leverages a first-order approximation property of well-trained policy networks, which is empirically verified to be accurate within a 2% error margin across various RL environments. The resulting algorithm, PolicyGradEx, efficiently estimates an aggregate task-affinity score matrix given a policy evaluation algorithm. Based on the estimated affinity score matrix, we cluster the $n$ objectives into $k$ groups by maximizing the intra-cluster affinity scores. Experiments on three robotic control and the Meta-World benchmarks demonstrate that our approach outperforms state-of-the-art baselines by 16% on average, while delivering up to $26\times$ faster speedup relative to performing full training to obtain the clusters. Ablation studies validate each component of our approach. For instance, compared with random grouping and gradient-similarity-based grouping, our loss-based clustering yields an improvement of 19%. Finally, we analyze the generalization error of policy networks by measuring the Hessian trace of the loss surface, which gives non-vacuous measures relative to the observed generalization errors.

Figures (3)

• Figure 1: An overview of our approach. Left: Run multitask training on all the tasks, $T_1, T_2, \dots, T_n$, and obtain to a meta-initialization policy $\pi_{\theta^\star}$. Store the projected gradients of a surrogate loss with meta-policy $\theta^\star$ for every transition from $t = 1, 2, \dots, N$. Middle: Estimate policy adaptation performance on $m$ task subsets using projected gradients as features in logistic regression. Right: Compute an $n\times n$ task affinity score matrix based on the estimated loss values. Lastly, run a clustering algorithm to group similar objectives, resulting in $k$ subgroups $G_1, \dots, G_k$, each of which share the same policy within group.
• Figure 2: Illustrating the Hessian trace measurements and empirical generalization errors with respect to the policy network. Figure \ref{['fig_hessian_training_curve']}: Showing that the Hessian trace is comparable in scale to the observed generalization errors, tested on Meta-World and a control task. Figure \ref{['fig_hessian_vary_k']}: Showing that in a Meta-World environment, the generalization error reaches the highest when the subset size $\alpha = 3$, suggesting negative transfer among a small set of tasks. In the meta-RL control task, the generalization performance monotonically improves with the addition of more tasks in each group.
• Figure 3: We illustrate the meta-training curve and the corresponding meta task distribution of the navigation environment. The agent is meta-trained to navigate a training set of goals, and then tested on a distinct set of unseen test goals with a few adaptation steps.

Theorems & Definitions (10)

• Theorem 1
• Theorem 2
• Proposition 3
• Claim 4
• proof
• Lemma 5
• proof
• Lemma 6
• proof : Proof of Theorem \ref{['thm_hessian']}
• proof : Proof of Lemma \ref{['lemma_union_bound']}