Multi-Objective Reinforcement Learning Using Routed Ensembles and Trajectory Attribution

AI summary

Overview

Research area: Multi-objective reinforcement learning (MORL), spanning regret theory, policy-learning algorithm design, and data-influence/attribution methods. The work sits at the intersection of RL theory, transfer learning under distribution shift, and data selection for post-training.

Technical level: Advanced. The paper combines a chi-squared-divergence-based conflict measure, upper and lower regret bounds, an online ensemble algorithm, and a kernel ridge regression surrogate for trajectory influence.

Scope (one sentence): The paper quantifies how misalignment among objectives governs the sample complexity of online MORL, then proposes a routed ensemble learner and a kernel-surrogate trajectory attribution method, validated on synthetic tasks, Meta-World, MO-Gymnasium, and an RLHF setting with Qwen3-1.7B.

What This Paper Is About

When a reinforcement learning agent must trade off several competing objectives (for example helpfulness versus honesty, or speed versus energy), the best policy can change sharply as the preference weighting changes, so training a separate policy for every weighting is sample-inefficient. The paper asks how much the degree of conflict between objectives drives that difficulty, and whether a small set of reusable "basis" policies plus a routing rule can cover many preferences at lower cost. It also asks a complementary question: which individual training trajectories actually help or hurt policy learning, and can that influence be estimated cheaply and accurately enough to select better data.

Key Contributions

1. A theory of objective conflict. The authors introduce an objective conflict coefficient Γ, defined as the maximum chi-squared divergence between each basis policy's state-action occupancy and the uniform average of the K basis occupancies. They derive an upper regret bound for a K-policy optimistic confidence-set algorithm and a nearly matching lower bound, making the dependence on misalignment explicit.

2. Routed ensemble policy learning (Algorithm 1). An online algorithm that maintains N basis policies, scores each incoming trajectory under every basis policy by its estimated scalarized return, routes the trajectory to the highest-scoring policy, and updates that policy only on its assigned data using a standard RL subroutine. Lemma 3.8 proves that routing to a single best basis policy is optimal when the composite policy lies in the convex hull of the bases.

3. Kernel-surrogate trajectory attribution (Algorithm 2). A scalable alternative to leave-one-out (LOO) retraining that samples n trajectory subsets, estimates each subset's performance locally around a reference policy using policy gradients from local rollouts, and fits kernel ridge regression with a radial basis function kernel exp(−γ‖Sᵢ − Sⱼ‖²). The nonlinear kernel captures synergy and antagonism between trajectories, which first-order influence approximations miss. The same procedure can be applied to objective attribution by sampling subsets of objectives.

4. Broad empirical validation. Results on synthetic environments (a 4-objective Pareto-frontier task and a 2-objective corrupted-signal task), Meta-World MT10, MO-Gymnasium continuous and discrete tasks, Multitask Minigrid, a synthetic synergy MTRL environment, and a HelpSteer2 RLHF setting using Qwen3-1.7B fine-tuned with PPO via per-basis LoRA adapters.

Main Findings

• Conflict enters the regret bound explicitly. Theorem 3.6 gives Reg(T) ≤ O_T(H√(C_cov(1+Γ) log(K|F|Tδ⁻¹)T)), where C_cov is the coverability coefficient and K the number of basis policies; the burn-in is t₀ = O(C_cov K β) with β = c₁ log(K|F|T/δ). A matching lower bound of order Ω_T(H√((1+Γ)T)) in Theorem B.6 establishes tightness in Γ and T up to C_cov and logarithmic factors.

• Γ is bounded by the number of bases. Lemma B.5 shows Γ ≤ K − 1, so the conflict measure scales with ensemble size; Γ = 0 in the aligned regime and is larger in the strictly concave frontier case than in the linear case.

• Routing relieves the exploration bottleneck. After the routed ensemble procedure, the coverability metric drops from 1.3 for a single policy to 0.2 for the ensemble policies.

• Reported gains over baselines. The abstract states a 4.3% average improvement over baselines for online environments from Meta-World and MO-Gymnasium. Section 1 reports 4.3% over the strongest baseline on Meta-World and an 8% average hypervolume improvement on MO-Gymnasium. Section 4 reports a 4.3% relative margin in multi-task RL, up to 48% in multi-objective RL, and 17% on large language model benchmarks.

• Attribution quality. Kernel surrogate attribution improves Spearman correlation with ground-truth leave-one-out influence scores by 50% over existing baselines, and Algorithm 2 achieves a Spearman rank correlation of 0.75 with true policy-gradient influence scores.

• Better data selection. When used for trajectory selection, kernel-based attribution yields a 27% improvement in reward over baseline attribution methods.

• Synthetic validation. With 4 conflicting objectives, routed ensembles substantially reduce cumulative regret relative to MO-UCBVI (which learns one policy per objective); under a corrupted second reward with target preference [1, 0], routing mitigates negative transfer and reaches near-optimal performance.

• The low-rank assumption holds approximately in practice. The mean convex-hull residual falls from 0.186 at K = 3 to 1.45 × 10⁻⁴ at K = 21 in a tabular concave setting, and from 0.443 at K = 3 to 0.231 at K = 15 on MO-Hopper using candidate-selected bases.

• Attribution tracks conflict. The fitted kernel model evaluated at [1]_T − e_t is strongly positively correlated with LOO scores and strongly negatively correlated with the objective conflict coefficient Γ (Figure 5, Appendix C.2).

Methodology in Plain English

The environment is modeled as a multi-objective Markov decision process with vector rewards r ∈ ℝᵐ and a preference vector w; performance under a preference is the linear scalarization ⟨w, J(π)⟩ averaged over a preference distribution. The theory assumes that optimal policies across all preferences can be expressed as mixtures of K latent basis policies (an exact convex-hull condition, relaxed in Appendix B.4 to an ε-neighborhood), and that preferences are sampled broadly enough to keep every basis exercised.

Conflict is measured by comparing each basis policy's visitation pattern to the average visitation pattern of all bases, using chi-squared divergence. If all bases visit states similarly, the coefficient is small and information transfers well across preferences; if they visit nearly disjoint states, it is large. The regret analysis runs a K-policy version of an optimistic confidence-set algorithm, maintaining one confidence set per basis and taking a union bound across them.

The practical algorithm replaces latent bases with N learned policies. Each round it samples a preference, rolls out B trajectories from randomly chosen basis policies, scores each trajectory against every basis using the estimated scalarized return, assigns it to the argmax policy, and updates each policy only on its own cluster. Isolating updates matters because a policy tuned for one preference can degrade under another.

For attribution, exact leave-one-out retraining is too expensive and captures only one-at-a-time effects. Instead the authors sample subsets of trajectories, estimate each subset's value with a local policy-gradient approximation around a policy trained on all data (avoiding full retraining in the online setting), and fit a kernel ridge regression over subset indicator vectors. The RBF kernel lets the surrogate model interactions rather than assuming additive contributions. Trajectories with the highest attributed scores are then selected for downstream training.

Why This Matters

Impact on research. The paper is the first, by its own claim, to formalize how objective misalignment affects learning efficiency across preferences, tying a divergence-based conflict measure to the established coverability complexity measure. It generalizes single-objective regret results and supplies a matching lower bound, and it reframes data influence estimation as a nonlinear, kernel-based problem rather than a first-order one.

Real-world applications:

• Robotics: balancing speed against energy consumption, and multi-task manipulation where tasks share dynamics (the MT10 setup).
• LLM post-training and RLHF: handling competing criteria such as helpfulness, safety, and honesty, and detecting corrupted or adversarial reward components.
• Data curation: scoring and filtering training trajectories before policy learning, which the paper shows improves downstream reward by 27% over baseline attribution methods.
• Multi-objective control benchmarks: continuous locomotion and discrete resource-trading tasks where explicit Pareto trade-offs exist.

Industry relevance. Routed ensembles are applied as a drop-in module on top of published base learners, so teams can add them without replacing existing training stacks. Avoiding per-preference retraining and full leave-one-out analysis directly reduces compute, which matters for expensive post-training pipelines and for large-scale trajectory datasets.

Future Directions

• Choosing and growing the ensemble size. The paper assumes N is given and suggests cross-validation (Section 4.3) or adaptively spawning new clusters when existing policies fail to explain incoming trajectories, drawing on clustering literature.

• Further relaxing the low-rank basis assumption. Appendix B.4 treats approximate low-rankness, and Proposition B.10 preserves the qualitative dependence on K and Γ, but the empirical residuals on MO-Hopper remain as large as 0.231 at K = 15, leaving room for tighter approximations.

• More accurate and scalable subset-performance estimation. The attribution pipeline relies on a local policy-gradient surrogate rather than retraining; better estimators could improve fidelity on larger trajectory sets and in non-stationary online data.

• Extending the theory beyond the stylized setting. Theorem 3.6 targets a K-policy optimistic algorithm with linear scalarization; extending guarantees to function approximation, other scalarization operators, and the objective-attribution variant is left open.

Target Audience

MORL and RL theory researchers interested in sample complexity under competing objectives; RLHF and post-training engineers who need to balance multiple reward models and curate training data; robotics practitioners working on multi-task control; and data-influence or data-selection researchers who want a nonlinear alternative to influence functions and linear surrogates. Readers should be comfortable with occupancy measures, regret analysis, and kernel methods.

Authors’ abstract

Multi-objective reinforcement learning is a fundamental problem for modern AI, with applications spanning post-training, reasoning models, and robotic control. A central challenge is to learn policies that balance competing objectives, such as helpfulness versus honesty in safety and alignment, or speed versus energy in robotics. The optimal policy can vary sharply with objective-weighting combinations, and the extent of this variation depends crucially on the degree of conflict among objectives, making per-preference policy learning sample-inefficient. In this paper, we introduce objective-conflict coefficients and derive near-optimal regret bounds for online settings using this measure, thereby quantifying the impact of misalignment on learning efficiency. Then, we design routed ensemble policy learning, an algorithm that maintains a small set of basis policies and dynamically routes trajectories to the policy with the highest potential reward. Finally, we consider trajectory selection as another approach to mitigate objective conflict, and design a kernel surrogate trajectory attribution algorithm that estimates how feedback signals from a set of trajectories influence policy learning on other trajectories. The kernel can capture nonlinear interactions such as synergy and antagonism between trajectories. We evaluate our approach across diverse environments and benchmarks: routed ensembles improve average performance by 4.3% over baselines for online environments from Meta-World and MO-Gymnasium, and kernel surrogate attribution improves correlation with leave-one-out analysis over existing baselines by 50%. When applied to trajectory selection, kernel-based attribution yields a 27% improvement in reward over baseline attribution methods.