RILe: Reinforced Imitation Learning

TL;DR

RILe tackles reward engineering in high dimensional control by coupling a trainer that learns a dense reward with a student that imitates expert behavior in a unified RL framework. The trainer uses a discriminator guided signal and a novel reward $R^T = e^{|\upsilon(D_\phi(s^T)) - a^T|}$ to provide context sensitive feedback as the student evolves, enabling on the fly reward shaping. Empirically, RILe outperforms state of the art adversarial IL and IRL methods on MuJoCo and LocoMuJoCo benchmarks, achieving near expert performance in several tasks while reducing the computational burden of traditional IRL cycles. The approach demonstrates strong robustness to noise and covariate shift and reveals important trade offs in expert data usage, offering a practical, dynamic alternative to static reward learning for complex robotic and control problems.

Abstract

Acquiring complex behaviors is essential for artificially intelligent agents, yet learning these behaviors in high-dimensional settings poses a significant challenge due to the vast search space. Traditional reinforcement learning (RL) requires extensive manual effort for reward function engineering. Inverse reinforcement learning (IRL) uncovers reward functions from expert demonstrations but relies on an iterative process that is often computationally expensive. Imitation learning (IL) provides a more efficient alternative by directly comparing an agent's actions to expert demonstrations; however, in high-dimensional environments, such direct comparisons often offer insufficient feedback for effective learning. We introduce RILe (Reinforced Imitation Learning), a framework that combines the strengths of imitation learning and inverse reinforcement learning to learn a dense reward function efficiently and achieve strong performance in high-dimensional tasks. RILe employs a novel trainer-student framework: the trainer learns an adaptive reward function, and the student uses this reward signal to imitate expert behaviors. By dynamically adjusting its guidance as the student evolves, the trainer provides nuanced feedback across different phases of learning. Our framework produces high-performing policies in high-dimensional tasks where direct imitation fails to replicate complex behaviors. We validate RILe in challenging robotic locomotion tasks, demonstrating that it significantly outperforms existing methods and achieves near-expert performance across multiple settings.

Figures (6)

• Figure 1: Overview of the related works.(a) Reinforcement Learning (RL): learning a policy that maximizes hand-defined reward function; (b) Inverse RL (IRL): learning a reward function from data. IRL has two stages: 1. training a policy with frozen reward function, and 2. updating the reward function by comparing the converged policy with data. These stages repeated several times; (C) Generative Adversarial Imitation Learning (GAIL) + Adversarial IRL (AIRL): using discriminator as a reward function. GAIL trains both policy and the discriminator at the same time. AIRL implements a new structure on the discriminator, separating reward from environment dynamics by using two networks under the discriminator (see additional terms in green). (D) RILe: similar to IRL, learning a reward function from data. RILe learns the reward function at the same time with the policy, using a discriminator as a guide for learning the reward.
• Figure 2: Reinforced Imitation Learning (RILe). The framework consists of three key components: a student agent, a trainer agent, and a discriminator. The student agent learns a policy $\pi_S$ by interacting with an environment, and the trainer agent learns a reward function as a policy $\pi_T$. (1) The student receives the environment state $s^S$. (2) The student takes an action $a^S$, forwards it to the environment which is updated based on $a^S$. (3) The student forwards its state and action to the trainer, whose state is $s^T = (s^S, a^S)$. (4) Trainer, $\pi_T$, evaluates the state action pair of the student agent $s^T = (s^S, a^S)$ and chooses an action $a^T$ that then becomes the reward of the student agent $a^T = r^S$. (5) The trainer agent forwards the $s^T = (s^S, a^S)$ to the discriminator. (6) Discriminator compares student state-action pair with expert demonstrations ($s^D$). (7) Discriminator gives reward to the trainer, based on the similarity between student- and expert-behavior.
• Figure 3: Reward Function Comparison. Evolution of reward functions during training for (a) RILe, (b) GAIL, and (c) AIRL in a continuous maze environment. Columns show reward landscapes at 25%, 50%, 75%, and 100% of training completion (left to right). The expert's trajectory is shown in red, while the student agent's trajectory from the previous training epoch is in black. Color gradients represent reward values, with darker colors indicating lower rewards and brighter colors indicating higher rewards. Black squares represent obstacles. RILe demonstrates a dynamic reward function that adapts with the student's progress, while GAIL and AIRL maintain relatively static reward landscapes throughout training and struggle to adapt.
• Figure 4: Dynamics of Reward Functions. (a) Reward Function Distribution Change (RFDC): Wasserstein distance between reward function distributions. (b) Fixed-State Reward Function Distribution Change (FS-RFDC): Mean absolute deviation of reward values for a fixed set of expert states. (c) Correlation between Performance and Reward (CPR): Pearson correlation between changes in the reward function and changes in the student's performance.
• Figure 5: Trainer-Discriminator Relation: Comparison of different trainer reward functions, each defining a different relationship between the trainer’s action and the discriminator’s output. The student’s return curves on the left show how performance evolves, and the normalized final performance on the right presents a clear comparison between reward designs. Exponential naive converges faster but plateaus at a lower final reward, whereas exponential difference yields the highest performance.