PaperlayerView exact revision on arXiv

Source-linked AI summary

Scalable Multi-Task Inverse Reinforcement Learning

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

arXiv:2610.00758v1cs.LGstat.ML

TL;DR

Counterfactual reward transfer is difficult when each task’s demonstrations cover only part of the state space. The paper pools tasks through a low-rank policy-logit model and shared reward bases; experiments report stronger off-support recovery and target transfer alongside computation that scales with rank rather than task count.

  • Problem

    In heterogeneous tasks, limited state coverage weakens reward identification, even though transferable rewards are needed to predict behavior under changed environments.

    Source: §1 · p. 1

  • Method

    The method jointly fits low-rank source-policy logits across tasks and uses shared basis functions to recover distinct rewards with planning effort tied to rank.

    Source: §1 · p. 2

  • Results

    Across experiments, LowRank improves off-support reward recovery and counterfactual transfer, while its computation grows as N^0.21 and reaches up to 13× faster than per-task GenPQR at N = 128.

    Source: §1 · p. 2, §5.3 · p. 8, §5.4 · p. 9

  • Takeaways & Limitations

    Pooling can recover reward variation beyond individual task support, while shared bases reduce planning from per-task to per-basis.

    Source: §6 · p. 9

    The finite-sample guarantees are limited to tabular MDPs with a known source kernel.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

IRL can use recovered rewards to predict how agents re-plan when environments change, but limited state coverage makes reward transfer difficult. This paper pools demonstrations across tasks under a low-rank logit assumption to reduce coverage demands and scale reward recovery with the shared rank.

1 INTRODUCTION

  • Single-task IRL faces a coverage–choice-sharpness tradeoff: high temperature broadens coverage but flattens choices, while low temperature sharpens rewards but narrows coverage.Fitting tasks independently inherits the single-task support problem.Source: §1 · p. 1
  • The method assumes low-rank policy logits, a structure motivated by shared dynamics and subgoals and made empirically testable through choice probabilities.The paper places its assumption on logits rather than rewards because rewards are affine in log policies in the analyzed maximum-entropy IRL representation.Source: §1 · p. 2
  • IRL’s expensive planning step scales with the number of shared basis functions k rather than the number of tasks N.The paper linearly composes k basis fixed points to obtain N task rewards, yielding a speed-up when k ≪ N.Source: §1 · p. 2
  • Pooling tasks under a low-rank assumption lets agents’ demonstrations jointly cover states, reducing the need for each task to visit every state.The method retains distinct rewards for different tasks while using shared information across their observed states.Source: §1 · pp. 1–2
  • The paper introduces an offline, feature-free method, proves finite-sample guarantees, and reports stronger off-support recovery, counterfactual transfer, and rank-scaled computation.Experiments span three domains and report wall-clock growth as N^0.21, up to 13× faster than per-task GenPQR at N = 128.Source: §1 · p. 2

2 RELATED LITERATURE

Related work spans single-task IRL, multi-task methods, and low-rank reward learning. This paper distinguishes its offline recovery of distinct task rewards from prior approaches that require interaction, known reward features, or per-task planning.

Single Task IRL

  • Single-task IRL includes maximum-entropy likelihood fitting, adversarial occupancy matching and soft Q-functions, while partial coverage limits offline reward identification.The paper also builds on the classify-then-regress view of van der Laan et al. (2025) and dynamic discrete-choice work on heterogeneous agent types.Source: §2 · p. 2

Multi-task IRL

  • Prior multi-task methods learn latent task embeddings or shared structure, while discrete multi-intention methods cluster rewards; this paper instead recovers N distinct rewards with continuous loadings on k bases.The cited continuous-latent and hierarchical methods require simulator interaction and adversarial optimization and do not provide the paper’s guarantees under missing coverage or rank-scaled inversion.Source: §2 · p. 2

Low rank rewards and matrix completion

  • Unlike low-rank reward methods for time-varying rewards or forward RL, this work imposes low rank on log policies and recovers rewards from offline behavior without reward features and with per-basis planning.Its analysis allows unobserved task–state blocks through a restricted coverage constant and carries completion error through the Bellman reward operator to target-policy KL.Source: §2 · p. 2

3 METHOD

The method fits shared low-rank source-policy logits from demonstrations, then recovers task rewards through shared basis functions. Its practical basis-planning shortcut reduces inversion work, while its theory distinguishes that computational approximation from the exact normalized estimator.

3.1 SETUP

  • Source and target MDPs share states, actions, rewards, and the initial-state distribution, but may differ in transition kernels and discount factors.The theoretical analysis assumes known dynamics; practical source reward recovery uses observed transitions, while target re-planning requires known target dynamics.Source: §3.1 · p. 3

3.2 MULTI-TASK POLICY LEARNING

  • The method models source policy logits as a rank-at-most-k product of task-specific loadings and shared basis functions, then jointly estimates the factors by penalized maximum likelihood.A reference action is fixed to zero because softmax is invariant to a common shift of action logits.Source: §3.2 · p. 3
  • Shared bases let tasks pool information across observed states, so individual tasks need not cover every state if pooled observations distinguish errors compatible with the low-rank structure.The method learns rewards from the offline dataset without the true source transition kernel or reward features.Source: §3.2 · p. 3

3.3 EFFICIENT REWARD RECOVERY THROUGH SHARED BASES

  • Reward identification uses a normalization shared across tasks, and the low-rank logit model is a structural approximation because the nonlinear maximum-entropy optimality map can increase effective rank.The exact normalized plug-in reward includes a task-dependent log-partition normalizer; applying the reward operator to it can require one inversion per task.Source: §3.3 · pp. 3–4
  • The practical method applies the reward operator to anchored low-rank logits and recombines k basis rewards, rather than performing N task-specific inversions.The shortcut’s approximation gap is Lµbci, which need not vanish under action-dependent transitions; experiments report mild to null recovery effects and wall-clock savings that grow with N.Source: §3.3 · p. 4

3.4 ALGORITHM

  • Algorithm 1 jointly fits rank-k anchored logits, runs k fitted-Q-iteration regressions on learned bases, and linearly recombines the resulting rewards.A new task adds loadings and classification data, not another reward-inversion solve; continuous-state basis functions use neural networks.Source: §3.4 · p. 4, §3.3 · p. 4

4 THEORETICAL GUARANTEES

The theoretical analysis estimates shared low-rank source policy logits and uses pooled coverage to recover distinct task rewards despite missing task–state observations. Under its assumptions, it provides finite-sample guarantees for estimation and target-policy performance.

4 THEORETICAL GUARANTEES

  • The algorithm regresses k basis rewards and forms each task’s reward as a coefficient-weighted combination plus an offset.Source: §4 · p. 5

4.1 STATISTICAL ASSUMPTIONS AND ESTIMATOR

  • Under bounded rank-k logits, the estimator fits the source logit matrix by nuclear-norm-regularized maximum likelihood on pooled demonstrations.Source: §4.1 · p. 5

4.2 POOLED SOURCE COVERAGE

  • Pooled coverage can permit missing task–state blocks when the tasks visiting each state collectively span the latent directions; otherwise the coverage constant can be zero.An example shows full reward recovery and vanishing target regret can remain possible despite an unobserved task–state block.Source: §4.2 · p. 6

4.3 TARGET VALUES AND FINITE-SAMPLE GUARANTEES

  • With probability at least 1−δ, the theorem gives simultaneous finite-sample guarantees for estimation, target values, and policy performance across tasks and stationary target policies.The policy guarantee is expressed through occupancy-weighted KL divergence from each task’s soft-optimal target policy.Source: §4 · pp. 4–5, §4.3 · p. 6

5 EXPERIMENTS

Across three domains, the experiments test whether pooled coverage improves off-support reward recovery, whether recovered rewards transfer under target shifts, and whether computation scales sublinearly with task population. LowRank performs strongly across these evaluations, with larger computational advantages as tasks grow.

5 EXPERIMENTS

  • Experiments span FourRooms, Highway, and RecSim; for continuous domains, off-support states are observed only for other tasks, not arbitrary unseen points.Source: §5 · p. 7

5.2 SAMPLE EFFICIENT REWARD RECOVERY

  • LowRank achieves the highest off-support reward correlation at every data budget across all three domains, with the largest gains at low data; per-task methods approach it only on Highway at 10,000 decisions.Source: Figure 2 · p. 8, §5.2 · p. 8

5.1 COVERAGE DRIVES OFF-SUPPORT RECOVERY

  • Targeted projected error grows with deletion dose, and the difference-in-differences is concentrated in the targeted direction, separating directional coverage loss from sample-volume loss.Source: §5.1 · pp. 7–8

5.3 POLICY RECOVERY IN PERTURBED ENVIRONMENTS

  • Across target perturbations, LowRank has the highest normalized target return; its slope is +2.6 percentage points per 0.1 off-support mass (95% CI [−3.2, 8.3]).Source: §5.3 · p. 8, Figure 3 · p. 9

5.4 COMPUTATIONAL EFFICIENCY

  • At N=128, LowRank is roughly 13×, 6×, and 4× faster than per-task GenPQR on Highway, FourRooms, and RecSim, respectively.Source: Figure 4 · p. 9, §5.4 · p. 9

6 CONCLUSION

The method pools anchored source logits to relax per-task coverage requirements and reduce planning to per-basis computation, while its guarantees and validation remain bounded by stated assumptions and settings.

6 CONCLUSION

  • Identifiability does not require full per-task state coverage; demonstrations must instead span the low-rank cone under a restricted-eigenvalue condition.At minimum, tasks visiting each state must span latent task directions, and each task must visit states spanning latent basis directions.Source: §6 · p. 9
  • Pooling recovers the correct direction of reward variation outside task support and yields sub-linear wall-clock scaling as task count grows.Planning is reduced from per-task to per-basis.Source: §6 · p. 9
  • Finite-sample guarantees are limited to tabular MDPs with a known source kernel; extending theory to approximation error and relaxing shared reward normalization remain open.The environments are synthetic, and validation on real behavioral data remains open.Source: §6 · pp. 9–10

AI USE STATEMENT

This block records the authors’ AI use, appendix notation and setting assumptions, and the procedure for fitting source logits and selecting target policies.

AI USE STATEMENT

  • Generative AI assisted with experiment coding, figure generation, and checks for notation mismatches and typos; authors reviewed the work and take responsibility for the final content.Source: AI USE STATEMENT · p. 10

A THEORY APPENDIX

  • The appendix defines matrix norms and inner products, vector norms, and task-state action blocks used in its notation.Source: A THEORY APPENDIX · p. 12

A.1 SETTING

  • Tasks are finite MDPs with shared mean rewards across source and target environments; transitions, initial distribution, and discounts are known, and demonstrations sample source soft-optimal actions.Source: §A.1 · pp. 12–13

A.2 PROCEDURE: SINGLE LOGIT FIT

  • The procedure fits source logits, reconstructs each task’s reward, estimates candidate-policy target soft values, then selects the highest-valued policy.Source: §A.2 · pp. 14–15

A.3 MAIN RESULT: SINGLE LOGIT FIT

This section analyzes single-logit fitting, giving guarantees for estimation and for transferring fitted-logit error to target-policy values and regret under stated coverage and normalization conditions. It also addresses reward recovery and exact recovery when the true parameter is zero, and describes the experimental setup and baselines.

A.3 MAIN RESULT: SINGLE LOGIT FIT

  • With probability at least 1 −δ, Theorem 2 simultaneously guarantees common-logit estimation and estimated target values across tasks and candidate policies.The guarantee applies to every minimizer in (28) and every maximizing choice in (32); the supplied passages omit the bound expressions.Source: §A.3 · p. 15

A.4 COVERAGE CONDITION

  • Assumption 6 requires a positive coverage coefficient c0, whose magnitude depends on the source design and the true logit matrix's low-rank geometry.The coverage coefficient is introduced for the coverage condition.Source: §A.4 · p. 16

A.4.1 INTERPRETING THE COVERAGE CONDITION

  • The coverage coefficient is positive exactly when no nonzero error in the anchored cone vanishes on all observed blocks.Equivalently, ker(Pd) ∩ C(Θ⋆) = {0}; positive restricted coverage also supports identification within the error cone.Source: §A.4.1 · pp. 18–19

A.5 ESTIMATION FOR THE SINGLE LOGIT FIT

  • Under the source sampling model and Assumptions 4 and 6, Lemma 8 gives a high-probability logit-estimation guarantee for every minimizer of the single-fit estimator.The supplied passage states the guarantee's conditions and scope but omits its bound expression.Source: §A.5 · p. 20
  • Population restricted strong convexity on feasible cone directions combines with empirical curvature beyond t0 to support the estimation bound.The population loss has the stated RSC, and the empirical curvature result applies to feasible cone directions with Frobenius norm at least t0.Source: §A.5 · pp. 24, 30
  • On the joint good and fixed-radius events, every fitted minimizer lies within deterministic radius t0 of the true logits, yielding the single-logit estimation bound.The events hold simultaneously with probability at least 1 − δ.Source: §A.5 · p. 31

A.6 REWARD IDENTIFICATION AND VALUE TRANSFER

  • Source-logit error propagates into simultaneous target-policy value-error bounds, which plug-in selection uses to bound regret by comparator and selected-policy radii.The value-error bounds cover all tasks and candidate policies.Source: §A.6 · p. 32

A.6.1 REWARD RECOVERY UNDER THE NORMALIZATION

  • Under Assumption 5, each task’s source soft-optimal log-policy uniquely determines its reward, and recovery is affine for fixed source dynamics, discount, and normalization.Source: §A.6.1 · p. 32

A.6.2 FROM FITTED LOGITS TO TARGET-POLICY VALUES

  • Under Assumption 5 and Θ⋆ ∈ MB, a fitted-logit error bound implies simultaneous value-error bounds for every task and candidate policy.The result applies to any feasible fitted matrix and the corresponding constructed estimates; the supplied passages do not show the bound’s formula.Source: §A.6.2 · p. 33

A.7.1 SELECTION FROM SIMULTANEOUS VALUE CERTIFICATES

  • Selecting the policy with the highest estimated value incurs loss bounded by the sum of the error radii for the comparator and selected policy.The selection guarantee is stated for a fixed task and estimates satisfying the simultaneous value certificates.Source: §A.7.1 · p. 35

A.7.2 SOFT REGRET AS KL DIVERGENCE

  • Soft regret for a stationary target policy is related to its occupancy-weighted KL divergence from the optimal target policy.Lemma 13 establishes this identity under the target dynamics and normalized occupancy definition.Source: §A.7.2 · p. 35

A.8 THE ZERO-TRUTH CASE

  • When Θ⋆ = 0, the estimator recovers zero with probability at least 1 − δ/2 without restricted coverage, and under Assumption 5 its value estimates are exact for every task and candidate policy.The exact value equality follows from recovery of the true rewards.Source: §A.8 · p. 37

B.1 ENVIRONMENTS

  • Experiments cover FourRooms, Highway, and RecSim, with sparse collection rewards, stochastic traffic rewards, and topic-click rewards, respectively.The settings use 8, 5, and 8 tasks; experts are trained with PPO in FourRooms and Highway and exact soft value iteration in RecSim.Source: §B.1 · pp. 37–38

B.2 REWARD CONSTRUCTION AND EXPERT TRAINING

  • Across environments, task weights combine pure and mixed tasks spanning R^4, with k = 4 and ground-truth rewards µ-normalized using uniform µ and g = 0.Source: §B.2 · p. 38

B.3 BASELINES

  • Baselines include per-task and pooled variants of GenPQR and IQ-Learn, plus feature-free Multi-intent IRL with four soft-clustering clusters.Multi-intent IRL recovers cluster rewards and plans once per cluster.Source: §B.3 · p. 38

B.4 ARCHITECTURES, HYPERPARAMETERS, AND TUNING

Within each environment, methods share an encoder and width, so comparisons isolate estimation strategy rather than model capacity. Regularization is tuned with matched opportunities, with the classifier’s selected setting transferred across environments.

B.4 ARCHITECTURES, HYPERPARAMETERS, AND TUNING

  • The classifier’s c = 0.2 setting won at every Highway budget and transferred unchanged to FourRooms and RecSim; no tested GenPQR setting improved its c = 0 baseline.IQ-Learn instead uses a fixed soft-Q temperature τ = 1.0.Source: §B.4 · p. 38
  • All methods within an environment share an encoder and width, isolating estimation strategy rather than model capacity.FourRooms uses a CNN, while Highway and RecSim use MLPs.Source: §B.4 · p. 38, Table 1 · p. 39

B.5 METRICS

The paper evaluates reward recovery on and off task support and measures transfer returns against an oracle ceiling. It also specifies how reward weights are generated across agents and environments.

B.5 METRICS

  • Reward recovery is measured by Pearson correlation with µ-normalized ground-truth rewards, separately on task-visited states and the remaining states in the union of visited states.The correlation is invariant to the global scale and shift that IRL cannot identify.Source: §B.5 · p. 39
  • Transfer return is reported as a fraction of the target-dynamics oracle ceiling while suppression of each task’s preferred topic is varied.The random-uniform and behavior-cloning policies bound the baselines.Source: §B.5 · p. 39

B.6 COMPUTATIONAL EFFICIENCY

  • Reward weights are sampled independently per agent, using environment-specific distributions, including nonnegative mixtures over a fixed rank-4 basis in RecSim.Highway weights are unit-normalized with nonpositive collision weight; FourRooms weights lie in the positive orthant.Source: §B.6 · p. 39

B.7 COVERAGE MECHANISM: DERIVATION OF THE WEIGHTED-NORM DECOMPOSITION

The weighted-norm decomposition links state coverage to error along the true logit subspace: weakly covered directions can hide large errors relevant to off-support recovery.

B.7 COVERAGE MECHANISM: DERIVATION OF THE WEIGHTED-NORM DECOMPOSITION

  • The decomposition uses the true logits’ left singular vectors and each task’s state-visit weights to form a state-specific weighted norm.The appendix derives the eigenvalue decomposition of the weighted-norm penalty stated in Section 5.1.Source: §B.7 · p. 39
  • Off-support global Pearson r is mostly flat as population size N grows for LowRank and both per-task baselines; MI varies by environment and sample size.The largest change for LowRank and the per-task baselines is 0.048; MI gains +0.51 on Highway at n=100 but loses 0.08 on FourRooms at n=500.Source: Table 2 · p. 40
  • Small eigenvalues mark directions with low coverage, where large error components contribute little to the weighted penalty despite being large in the unweighted Frobenius norm.The error is projected onto the column space of the true logits, and each state's visit weights shape its contribution.Source: §B.7 · pp. 39–40

B.8 SCALING TASKS

The transfer variant learns a basis from Phase 1 tasks, then fits loadings for Phase 2 tasks with that basis frozen. Increasing the Phase 1 pool improves LowRank’s off-support correlation for new tasks, while MI declines.

B.8 SCALING TASKS

  • The transfer approach learns basis Φ on a smaller Phase 1 task pool, then fits Phase 2 task loadings W by cross-entropy with the basis frozen.The Phase 1 tasks should span the basis space and provide as much data as possible for identification and statistical efficiency.Source: §B.8 · pp. 40–41
  • Increasing Phase 1 tasks from 8 to 128 raises LowRank’s off-support Pearson correlation by ≈0.15 for new Phase 2 tasks, while MI falls by 0.25.Figure 5 varies the number of Phase 1 tasks M and Phase 2 tasks Np2.Source: Figure 5 · p. 41, §B.8 · p. 41
  • Phase 1 task curation matters: selected tasks should span the basis space and have maximal data for identification and statistical efficiency.Source: §B.8 · p. 41

B.9 LOGIT FINITE SAMPLE BOUND RATE SCALING

The appendix tests predicted finite-sample scaling, the raw-logit approximation, rank misspecification, and cluster-count scaling. Results show dimension-dependent scaling, mild approximation effects, graceful rank degradation, and weaker MI performance than LowRank even at its best cluster count.

B.9 LOGIT FINITE SAMPLE BOUND RATE SCALING

  • Empirical rates match theory for sample size n and state cardinality S, but are more conservative for action cardinality A and rank k.The scaling study varies one dimension at a time around a fixed reference configuration.Source: §B.9 · p. 41, Figure 6 · p. 42

B.10 RAW LOGIT APPROXIMATION

  • Five of 12 paired comparisons favor exact normalization, five favor the approximation, and two are indistinguishable from zero; exact planning is slower.Exact planning takes approximately 1.1× the approximation’s wall-clock on Highway, 1.9× on FourRooms, and 1.8× on RecSim.Source: §B.10 · p. 41, Figure 7 · p. 42

B.11 RANK MISSPECIFICATION

  • Under- and over-specifying rank yields smooth recovery: lower k degrades gracefully, while higher k plateaus in RecSim and performance keeps increasing in FourRooms and Highway.RecSim plateaus at k=4, matching the reward-feature rank; continuous states are associated with greater true log-policy rank in FourRooms and Highway.Source: Figure 8 · p. 43, §B.11 · p. 43

B.12 MULTI-INTENT SCALING NUMBER OF CLUSTERS

  • Even at its best KMI per cell, Multi-intent performance is lower than LowRank across domains and both sample budgets.The effect of increasing cluster count is mixed: matching clusters to tasks sometimes performs best, while intermediate K can win in other settings.Source: Table 3 · p. 43, §B.12 · p. 43

Related papers

Loading 2610.00758v1…