Scalable Multi-Task Inverse Reinforcement Learning

Allen Tran; Jia Wan; Nathan Kallus; Aurélien Bibaut

cs.LGSep 30, 2026

Scalable Multi-Task Inverse Reinforcement Learning

Authors: Allen Tran, Jia Wan, Nathan Kallus, Aurélien Bibaut

Organizations: Netflix · Massachusetts Institute of Technology · Cornell Tech

Abstract

By learning transferable rewards, inverse reinforcement learning (IRL) enables counterfactual evaluation of agents under modified environments. Such transfer places strict requirements on coverage since target environments affect agents' state occupancy. We propose a multi-task IRL method that pools data across multiple agents with different rewards in the same environment under a low-rank assumption. In addition to alleviating coverage requirements, so each task need not visit every state as long as others do, the method offers scalable evaluation of multiple tasks under new environments as computationally intensive planning scales with rank rather than the number of tasks. We provide finite sample guarantees on reward recovery and on policy learning in new environments. Experiments show our method is robust to limited coverage, recovers rewards on and off of each task's support, transfers to target environments at lower regret than baselines, with its computational advantage over per-task methods widening as tasks grow.

Figures & tables

Figure 1: Causal test of directional coverage. Left: change in error projected onto the targeted eigendirection for treated, volume-matched control, and placebo cells. Right: treated-minus-control difference-in-differences for the targeted and remaining directions.

Figure 2: Reward recovery versus samples per task (mean and standard deviation over 10 seeds). LowRank leads off-support recovery at every budget; per-task methods narrow the gap only on Highway at high sample sizes.

Figure 3: Policy recovery in target perturbed environments. Error bands are

\pm

1 s.d over ten seeds.

Figure 4: Wall-clock scaling with task population

N

at fixed basis rank

k

. LowRank exhibits substantially sublinear growth, while per-task methods scale approximately linearly

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

	FourRooms	Highway	RecSim
Encoder	CNN	MLP	MLP
Hidden width	512	128	256
Residual blocks	0	0	2
FQI iterations	16	16	4
FQI regression epochs	50	50	40
Basis rank $k$	4	4	4

Appendix

Table 1: IRL pipeline settings per environment. All methods within an environment use the same encoder and width, so the comparison isolates the estimation strategy.

Domain	Method	$N{=}8$	$N{=}32$	$N{=}128$
Highway, $n{=}500$	LowRank	$+0.912\pm 0.010$	$+0.885\pm 0.011$	$+0.884\pm 0.010$
	GenPQR-per-task	$+0.787\pm 0.039$	$+0.742\pm 0.052$	$+0.739\pm 0.026$
	IQ-per-task	$+0.677\pm 0.040$	$+0.649\pm 0.042$	$+0.645\pm 0.024$
	MI	$+0.849\pm 0.014$	$+0.846\pm 0.082$	$+0.878\pm 0.003$
Highway, $n{=}100$	LowRank	$+0.879\pm 0.018$	$+0.847\pm 0.017$	$+0.844\pm 0.016$
	GenPQR-per-task	$-0.109\pm 0.082$	$-0.113\pm 0.029$	$-0.097\pm 0.020$

Appendix

Table 2: Reward-recovery quality (off-support global Pearson

r

, mean

\pm

s.d. over 10 seeds) against population size

N

. LowRank and both per-task baselines are flat in

N

(largest change is

0.048

); the small step from

N{=}8

N{=}32

on Highway and FourRooms is taken by all methods alike and reflects those populations sampling reward weights per agent, so agents

8

–

31

differ from the first

8

. MI is the exception: it gains

+0.51

on Highway at

n{=}100

, where its clusters are starved at

N{=}8

(s.d.

0.204

, falling to

0.028

N{=}128

), but loses

0.08

on FourRooms at

n{=}500

with tight error bars.

Figure 5: Transfer learning is effective: training basis on

M

agents in Phase 1 and transferring to

N_{p2}

Phase 2 agents

Figure 6: Empirical scaling in RecSim matches the theory in

n

and

S

and is more conservative in

k

and

A

Figure 7: Difference between reward recovery between reward operator applied to exact log normalized policies versus approximation

Figure 8: Misspecification of rank leads to graceful degradation

	$n=100$						$n=10^{4}$
$K_{\mathrm{MI}}$	Highway		FourRooms		RecSim		Highway		FourRooms		RecSim
$1$ †	$0.04$ $\pm$	$0.03$	$0.01$ $\pm$	$0.04$	$0.01$ $\pm$	$0.01$	$0.09$ $\pm$	$0.01$	$0.15$ $\pm$	$0.00$	$-0.00$ $\pm$	$0.00$
$2$	$\mathbf{0.23}$ $\pm$	$\mathbf{0.13}$	$0.13$ $\pm$	$0.08$	$0.04$ $\pm$	$0.06$	$0.54$ $\pm$	$0.03$	$0.37$ $\pm$	$0.06$	$0.18$ $\pm$	$0.01$
$3$	$0.04$ $\pm$	$0.07$	$0.18$ $\pm$	$0.07$	$0.08$ $\pm$	$0.06$	$0.60$ $\pm$	$0.13$	$0.41$ $\pm$	$0.03$	$0.30$ $\pm$	$0.03$
$4$	$-0.11$ $\pm$	$0.06$	$0.22$ $\pm$	$0.05$	$0.05$ $\pm$	$0.04$	$0.68$ $\pm$	$0.13$	$0.43$ $\pm$	$0.04$	$\mathbf{0.34}$ $\pm$	$\mathbf{0.02}$
$6$	—		$0.27$ $\pm$	$0.04$	$0.10$ $\pm$	$0.03$	—		$0.47$ $\pm$	$0.03$	$0.32$ $\pm$	$0.01$

Appendix

Table 3: MI-Cluster across cluster count

K_{\mathrm{MI}}

(off-support Pearson

r

, mean

\pm

s.d. over 10 seeds);

N{=}5

for Highway,

8

otherwise. Bold marks each column’s best

K

. Ours beats MI at its best

K

in every domain at both budgets. Shapes are budget-dependent: Highway rises to

K{=}N

n{=}10^{4}

but falls to it at

n{=}100

, and RecSim’s peak at the generative rank

K{=}4

disappears.

Explore similar work

Generalize and Guide: Decomposing Rewards for Few-Shot Inverse Reinforcement Learning
Jul 20, 2026Ziyi Liu, Grace ZhangOffline Reinforcement Learning Reward Signal
University of Southern California
Trust Region Inverse Reinforcement Learning: Explicit Dual Ascent using Local Policy Updates
May 10, 2026Anish Diwan, Davide Tateo, Christopher E. Mower +3Offline Reinforcement Learning Reward Functions
Technical University of Darmstadt · Robotics Institute Germany (RIG) · Lund University +3
ConTraIRL: Factorized Contrastive Abstractions for Transferable IRL
Jun 2, 2026Yikang Gui, Bikramjit Banerjee, Prashant DoshiOffline Reinforcement Learning Contrastive Learning
School of Computing, University of Georgia · School of Computing Sciences & Computer Engineering, The University of Southern Mississippi

Jul 20, 2026, cs.LGJ/K move · Enter open · S save

Generalize and Guide: Decomposing Rewards for Few-Shot Inverse Reinforcement Learning

Ziyi Liu, Grace Zhang

University of Southern California

Offline Reinforcement Learning Reward Signal Guidance Proximity Operations

Open paper AskPDF arXiv

Inverse reinforcement learning (IRL) provides a powerful framework for learning from demonstrations. However, real-world tasks often exhibit substantial natural variations (e.g., picking up mugs with varying shapes), making it impractical to collect demonstrations that fully specify a new task under every possible scenario. In practice, while demonstrations for the target task are limited, it is often easier to obtain datasets of heterogeneous but related behaviors. This motivates the problem of few-shot IRL with multi-task demonstrations (FM-IRL), where an agent must learn a new task with substantial variations from only a limited number of target-task demonstrations, together with sufficient demonstrations of related tasks and online agent experience. To do so, we must both recover the expert distribution of the new task and provide guidance when the agent deviates from it. We introduce Multitask discriminator Proximity-Guided IRL (MPG), which learns two complementary reward components: (1) a generalizable discriminator that transfers shared structure across related tasks to identify expert behavior in a new task, and (2) a proximity function that measures how far a state deviates from expert behavior and provides corrective guidance during exploration. We demonstrate the effectiveness of our method on multiple challenging navigation and manipulation tasks under significant variations (e.g., object configurations, table layouts, and initial robot poses), achieving an average success rate of 81.2%, outperforming the strongest per-task baseline by an average of 24.7 percentage points.

Abstract

Explore similar work

Generalize and Guide: Decomposing Rewards for Few-Shot Inverse Reinforcement Learning

Trust Region Inverse Reinforcement Learning: Explicit Dual Ascent using Local Policy Updates

ConTraIRL: Factorized Contrastive Abstractions for Transferable IRL