Imitation Learning as Return Distribution Matching

• URL: http://arxiv.org/abs/2509.12026v1
• Date: Mon, 15 Sep 2025 15:08:04 GMT
• Title: Imitation Learning as Return Distribution Matching
• Authors: Filippo Lazzati, Alberto Maria Metelli,
• Abstract summary: We study the problem of training a risk-sensitive reinforcement learning (RL) agent through imitation learning (IL)<n>Our goal is not only to train an agent that matches the expert's expected return but also its risk attitude.<n>We show that RS-KT achieves substantially lower sample complexity than RS-BC by exploiting dynamics information.
• Score: 37.79354987519793
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We study the problem of training a risk-sensitive reinforcement learning (RL) agent through imitation learning (IL). Unlike standard IL, our goal is not only to train an agent that matches the expert's expected return (i.e., its average performance) but also its risk attitude (i.e., other features of the return distribution, such as variance). We propose a general formulation of the risk-sensitive IL problem in which the objective is to match the expert's return distribution in Wasserstein distance. We focus on the tabular setting and assume the expert's reward is known. After demonstrating the limited expressivity of Markovian policies for this task, we introduce an efficient and sufficiently expressive subclass of non-Markovian policies tailored to it. Building on this subclass, we develop two provably efficient algorithms, RS-BC and RS-KT, for solving the problem when the transition model is unknown and known, respectively. We show that RS-KT achieves substantially lower sample complexity than RS-BC by exploiting dynamics information. We further demonstrate the sample efficiency of return distribution matching in the setting where the expert's reward is unknown by designing an oracle-based variant of RS-KT. Finally, we complement our theoretical analysis of RS-KT and RS-BC with numerical simulations, highlighting both their sample efficiency and the advantages of non-Markovian policies over standard sample-efficient IL algorithms.

Related papers

• Towards Sample-Efficient and Stable Reinforcement Learning for LLM-based Recommendation [56.92367609590823]
  Long Chain-of-Thought (Long CoT) reasoning has shown promise in Large Language Models (LLMs)<n>We argue that Long CoT is inherently ill-suited for the sequential recommendation domain.<n>We propose RISER, a novel Reinforced Item Space Exploration framework for Recommendation.
  arXiv  Detail & Related papers  (2026-01-31T10:02:43Z)
• ArenaRL: Scaling RL for Open-Ended Agents via Tournament-based Relative Ranking [84.07076200941474]
  ArenaRL is a reinforcement learning paradigm that shifts from pointwise scalar scoring to intra-group relative ranking.<n>We construct an intra-group adversarial arena and devise a tournament-based ranking scheme to obtain stable advantage signals.<n>Experiments show that ArenaRL substantially outperforms standard RL baselines.
  arXiv  Detail & Related papers  (2026-01-10T08:43:07Z)
• Mixture of Ranks with Degradation-Aware Routing for One-Step Real-World Image Super-Resolution [76.66229730098759]
  In real-world image super-resolution (Real-ISR), existing approaches mainly rely on fine-tuning pre-trained diffusion models.<n>We propose a Mixture-of-Ranks (MoR) architecture for single-step image super-resolution.<n>We introduce a fine-grained expert partitioning strategy that treats each rank in LoRA as an independent expert.
  arXiv  Detail & Related papers  (2025-11-20T04:11:44Z)
• Learning from Expert Factors: Trajectory-level Reward Shaping for Formulaic Alpha Mining [5.560011325936085]
  Reinforcement learning has successfully automated the complex process of mining formulaic alpha factors, for creating interpretable and profitable investment strategies.<n>Existing methods are hampered by the sparse rewards given the underlying Markov Decision Process.<n>To address this, Trajectory-level Rewarding (TLRS), a novel reward shaping method, is proposed.
  arXiv  Detail & Related papers  (2025-07-27T13:14:48Z)
• Risk-Averse Certification of Bayesian Neural Networks [70.44969603471903]
  We propose a Risk-Averse Certification framework for Bayesian neural networks called RAC-BNN.<n>Our method leverages sampling and optimisation to compute a sound approximation of the output set of a BNN.<n>We validate RAC-BNN on a range of regression and classification benchmarks and compare its performance with a state-of-the-art method.
  arXiv  Detail & Related papers  (2024-11-29T14:22:51Z)
• A Thorough Performance Benchmarking on Lightweight Embedding-based Recommender Systems [67.52782366565658]
  State-of-the-art recommender systems (RSs) depend on categorical features, which ecoded by embedding vectors, resulting in excessively large embedding tables.<n>Despite the prosperity of lightweight embedding-based RSs, a wide diversity is seen in evaluation protocols.<n>This study investigates various LERS' performance, efficiency, and cross-task transferability via a thorough benchmarking process.
  arXiv  Detail & Related papers  (2024-06-25T07:45:00Z)
• Heterogeneous Learning Rate Scheduling for Neural Architecture Search on Long-Tailed Datasets [0.0]
  We propose a novel adaptive learning rate scheduling strategy tailored for the architecture parameters of DARTS. Our approach dynamically adjusts the learning rate of the architecture parameters based on the training epoch, preventing the disruption of well-trained representations.
  arXiv  Detail & Related papers  (2024-06-11T07:32:25Z)
• Provable Risk-Sensitive Distributional Reinforcement Learning with General Function Approximation [54.61816424792866]
  We introduce a general framework on Risk-Sensitive Distributional Reinforcement Learning (RS-DisRL), with static Lipschitz Risk Measures (LRM) and general function approximation. We design two innovative meta-algorithms: textttRS-DisRL-M, a model-based strategy for model-based function approximation, and textttRS-DisRL-V, a model-free approach for general value function approximation.
  arXiv  Detail & Related papers  (2024-02-28T08:43:18Z)
• Is Inverse Reinforcement Learning Harder than Standard Reinforcement Learning? A Theoretical Perspective [55.36819597141271]
  Inverse Reinforcement Learning (IRL) -- the problem of learning reward functions from demonstrations of an emphexpert policy -- plays a critical role in developing intelligent systems. This paper provides the first line of efficient IRL in vanilla offline and online settings using samples and runtime. As an application, we show that the learned rewards can emphtransfer to another target MDP with suitable guarantees.
  arXiv  Detail & Related papers  (2023-11-29T00:09:01Z)
• GEC: A Unified Framework for Interactive Decision Making in MDP, POMDP, and Beyond [101.5329678997916]
  We study sample efficient reinforcement learning (RL) under the general framework of interactive decision making. We propose a novel complexity measure, generalized eluder coefficient (GEC), which characterizes the fundamental tradeoff between exploration and exploitation. We show that RL problems with low GEC form a remarkably rich class, which subsumes low Bellman eluder dimension problems, bilinear class, low witness rank problems, PO-bilinear class, and generalized regular PSR.
  arXiv  Detail & Related papers  (2022-11-03T16:42:40Z)
• SENTINEL: Taming Uncertainty with Ensemble-based Distributional Reinforcement Learning [6.587644069410234]
  We consider risk-sensitive sequential decision-making in model-based reinforcement learning (RL) We introduce a novel quantification of risk, namely emphcomposite risk We experimentally verify that SENTINEL-K estimates the return distribution better, and while used with composite risk estimate, demonstrates better risk-sensitive performance than competing RL algorithms.
  arXiv  Detail & Related papers  (2021-02-22T14:45:39Z)
• RMIX: Learning Risk-Sensitive Policies for Cooperative Reinforcement Learning Agents [40.51184157538392]
  We propose a novel MARL method with the Conditional Value at Risk (CVaR) measure over the learned distributions of individuals' Q values. We show that our method significantly outperforms state-of-the-art methods challenging StarCraft II tasks.
  arXiv  Detail & Related papers  (2021-02-16T13:58:25Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.