Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Can Optimal Transport Improve Federated Inverse Reinforcement Learning?

David Millard, Ali Baheri

TL;DR

The paper addresses learning a shared reward function from multiple heterogeneous agents without sharing trajectories. It replaces standard parameter averaging in federated IRL with an entropically regularized Wasserstein barycenter that fuses locally learned MaxEnt IRL rewards, preserving geometric structure and privacy. The authors prove stability and parameter-error bounds, and demonstrate through grid-world and Gymnasium experiments that the barycentric fusion yields more stable and transferable rewards than averaging, especially under heterogeneity. This approach offers a principled, communication-efficient framework for cross-environment generalization in robotics, with clear scalability and discretization challenges for future work.

Abstract

In robotics and multi-agent systems, fleets of autonomous agents often operate in subtly different environments while pursuing a common high-level objective. Directly pooling their data to learn a shared reward function is typically impractical due to differences in dynamics, privacy constraints, and limited communication bandwidth. This paper introduces an optimal transport-based approach to federated inverse reinforcement learning (IRL). Each client first performs lightweight Maximum Entropy IRL locally, adhering to its computational and privacy limitations. The resulting reward functions are then fused via a Wasserstein barycenter, which considers their underlying geometric structure. We further prove that this barycentric fusion yields a more faithful global reward estimate than conventional parameter averaging methods in federated learning. Overall, this work provides a principled and communication-efficient framework for deriving a shared reward that generalizes across heterogeneous agents and environments.

Can Optimal Transport Improve Federated Inverse Reinforcement Learning?

TL;DR

The paper addresses learning a shared reward function from multiple heterogeneous agents without sharing trajectories. It replaces standard parameter averaging in federated IRL with an entropically regularized Wasserstein barycenter that fuses locally learned MaxEnt IRL rewards, preserving geometric structure and privacy. The authors prove stability and parameter-error bounds, and demonstrate through grid-world and Gymnasium experiments that the barycentric fusion yields more stable and transferable rewards than averaging, especially under heterogeneity. This approach offers a principled, communication-efficient framework for cross-environment generalization in robotics, with clear scalability and discretization challenges for future work.

Abstract

In robotics and multi-agent systems, fleets of autonomous agents often operate in subtly different environments while pursuing a common high-level objective. Directly pooling their data to learn a shared reward function is typically impractical due to differences in dynamics, privacy constraints, and limited communication bandwidth. This paper introduces an optimal transport-based approach to federated inverse reinforcement learning (IRL). Each client first performs lightweight Maximum Entropy IRL locally, adhering to its computational and privacy limitations. The resulting reward functions are then fused via a Wasserstein barycenter, which considers their underlying geometric structure. We further prove that this barycentric fusion yields a more faithful global reward estimate than conventional parameter averaging methods in federated learning. Overall, this work provides a principled and communication-efficient framework for deriving a shared reward that generalizes across heterogeneous agents and environments.
Paper Structure (25 sections, 2 theorems, 20 equations, 3 figures, 4 tables)

This paper contains 25 sections, 2 theorems, 20 equations, 3 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries & Problem Formulation
  4. Inverse Reinforcement Learning
  5. Wasserstein Barycenter
  6. Problem Formulation
  7. Methodology
  8. Client-side Learning
  9. Aggregation
  10. Parameter Recovery
  11. Theoretical Results
  12. Empirical Results
  13. Grid-World Navigation
  14. Gymnasium Benchmarks
  15. Ablation
  16. ...and 10 more sections

Key Result

theorem 1

Suppose Assumptions ass:shared-support and ass:features hold. Each agent $i\in\{1,\dots,K\}$ runs local MaxEnt IRL and outputs a reward vector $\hat{r}_i\in\mathbb{R}^n$. Fix a shift $\sigma>0$ so that all entries of $\hat{r}_i+\sigma\mathbf{1}$ and of a shared population reward $r^\star$ are positi with $Z_{\min}:=\min\,\!\bigl(\min_i Z(\hat{r}_i),\,Z(r^\star)\bigr)$. Let $\bar{p}$ be the $2$-Was

Figures (3)

  • Figure 1: Local vs. mean vs. barycentric rewards for three heterogeneous 5×5 grid-world clients. Local MaxEnt IRL shows spurious obstacle attraction; parameter averaging blurs semantics. Wasserstein barycentric fusion preserves geometry, suppresses artifacts, and recovers a smooth, goal-consistent field across clients.
  • Figure 2: Learning dynamics across environments. The left panel shows results for Pendulum-v1 and the right for CartPole-v1. Curves report mean $\pm$ std true environment return versus PPO training timesteps under parameter averaging and Wasserstein barycentric reward fusion.
  • Figure 3: Learning dynamics across higher-dimensional environments. Mean $\pm$ std true environment return versus PPO training timesteps under parameter averaging and Wasserstein barycentric reward fusion. These cases required coarse discretization due to the high dimensionality of the support.

Theorems & Definitions (4)

  • theorem 1: Stability and parameter-error bound
  • proof : Proof sketch
  • theorem 2: Policy--performance bound for barycentric fusion
  • proof