RILe: Reinforced Imitation Learning

## Abstract

Abstract Reinforcement Learning has achieved significant success in generating complex behavior but often requires extensive reward function engineering. Adversarial variants of Imitation Learning and Inverse Reinforcement Learning offer an alternative by learning policies from expert demonstrations via a discriminator. However, these methods struggle in complex tasks where randomly sampling expert-like behaviors is challenging. This limitation stems from their reliance on policy-agnostic discriminators, which provide insufficient guidance for agent improvement, especially as task complexity increases and expert behavior becomes more distinct. We introduce RILe (Reinforced Imitation Learning environment), a novel trainer-student system that learns a dynamic reward function based on the student’s performance and alignment with expert demonstrations. In RILe, the student learns an action policy while the trainer, using reinforcement learning, continuously updates itself via the discriminator’s feedback to optimize the alignment between the student and the expert. The trainer optimizes for long-term cumulative rewards from the discriminator, enabling it to provide nuanced feedback that accounts for the complexity of the task and the student’s current capabilities. This approach allows for greater exploration of agent actions by providing graduated feedback rather than binary expert/non-expert classifications. By reducing dependence on policy-agnostic discriminators, RILe enables better performance in complex settings where traditional methods falter, outperforming existing methods by 2x in complex simulated robot-locomotion tasks.

## 1 Introduction

Reinforcement Learning (RL) has emerged as a powerful framework for teaching agents to perform complex tasks. In recent years, deep reinforcement learning has demonstrated remarkable success in replicating sophisticated behaviors, including playing Atari games, chess, and Go [, ]. However, these achievements often come at a cost: the tedious and challenging process of designing reward functions, as predicting the policy outcome from a manually crafted reward function remains notoriously difficult.

To overcome the reward engineering problem, Imitation Learning (IL) leverages expert demonstrations to learn a policy. Since vast amounts of expert data are required to accurately learn expert behaviors, Adversarial Imitation Learning (AIL) approaches, such as GAIL [], have been proposed as data-efficient alternatives. AIL employs a discriminator to measure similarity between learned behavior and expert behavior, rewarding the agent accordingly. While computationally efficient, AIL methods suffer from a critical limitation: the policy-agnostic nature of their discriminators. The discriminator lacks any inherent incentive to guide the agent towards expert-like behavior, in contrast to engineered reward functions in RL. Consequently, AIL methods face challenges in complex tasks requiring extensive exploration to find optimal actions. For instance, in digital locomotion tasks, AIL methods often struggle to consistently replicate expert performance [].

Inverse Reinforcement Learning (IRL) is another approach to alleviate reward engineering. Unlike IL, which directly learns expert behavior, IRL seeks to infer the underlying reward function that motivates the agent to acquire expert behaviors. The reward function and the agent are trained iteratively, with updates to the reward function based on the agent’s behavior. This iterative process renders IRL computationally expensive []. Adversarial Inverse Reinforcement Learning (AIRL) [] attempts to address this inefficiency by introducing a discriminator that enables simultaneous learning of the policy and reward function. However, in AIRL, the reward function is tightly coupled to the discriminator, potentially limiting its ability to capture complex task structures or long-term dependencies and inheriting the limitations of a policy-agnostic discriminators. This highlights the need for a method that can learn a more flexible reward function without the computational overhead of traditional IRL methods.

To overcome these challenges and effectively learn behaviors in complex settings, we propose Reinforced Imitation Learning (RILe) (Fig. 1-(d)). RILe aims to combine the ability to learn a reward function that actively guides the agent to imitate expert behavior with the computational efficiency of adversarial frameworks. At the core of RILe is a novel trainer-student system designed to address the shortcomings of existing methods:

- •
A student agent that learns to replicate the expert’s policy via RL in the environment
- •
A trainer agent that learns a reward function via RL and guides the student agent during training

By integrating the trainer-student dynamic, RILe decouples reward learning from policy learning and the discriminator, allowing each component to specialize and thereby overcome the limitations of policy-agnostic discriminators. While RILe utilizes a discriminator similar to those in adversarial frameworks, its role is fundamentally redefined. In RILe, the discriminator’s primary function is to provide feedback to the trainer agent by distinguishing expert data from student roll-outs. This feedback serves as the reward signal for the trainer, not directly influencing the student agent. The trainer leverages the discriminator’s feedback to learn a reward function that effectively guides the student agent. This approach enables more nuanced reward shaping, particularly beneficial in tasks requiring complex decision-making and extensive exploration.

Our contributions are two-fold:

- 1.
Decoupled Reward-function Learning: We introduce a novel approach where the trainer agent learns the reward function independently from both the student agent and the discriminator. Unlike existing methods that derive rewards directly from discriminator outputs, our trainer agent uses reinforcement learning to optimize the reward function based on the feedback from the discriminator. By focusing on long-term reward maximization, RL enables the trainer to distill inconsistent feedback from the discriminator into meaningful rewards, leading to better student performance.
- 2.
Dynamic Reward Customization: Our trainer agent dynamically adjusts rewards based on the student agent’s progress, facilitating a better learning experience and enabling accurate imitation of expert behavior in complex settings. This adaptive approach allows for more gradual learning, particularly in tasks where the optimal behavior may change depending on the agent’s current capabilities.

We evaluate RILe against state-of-the-art methods in AIL, and AIRL, specifically GAIL [], and AIRL []. Our experiments span three distinct scenarios: (1) Tailoring a reward function dynamically in a discrete maze task, (2) Investigating the impact of expert data on the trainer-student dynamics in a humanoid locomotion task, and (3) Imitating motion-capture data in continuous control tasks. The results demonstrate RILe’s superior performance, especially in complex tasks, and its ability to learn an effective dynamic reward function where baseline methods fail.

## 2 Related Work

We review literature on learning from expert demonstrations, focusing on Imitation Learning (IL) and Inverse Reinforcement Learning (IRL), which form the conceptual foundation of RILe.

#### Imitation Learning

Early work introduced Behavioral Cloning (BC) [], which learns a policy congruent with expert demonstrations through supervised learning. DAgger [] introduces data aggregation. GAIL [] introduces adversarial methods, where a discriminator aims to discriminate expert demonstrations, while a generator tries to fool the discriminator. DQfD [] proposes two-stage approach with pre-training, and ValueDice [] uses a distribution-matching objective between policy and expert. Despite progress, IL faces challenges in efficacy and generalization [, ]. RILe addresses these by introducing an adaptive teacher agent to guide the student beyond expert demonstrations.

#### Inverse Reinforcement Learning

IRL, introduced by [], learns the expert’s intrinsic reward function. Key developments include Apprenticeship Learning [], Maximum Entropy IRL [], and adversarial approaches like AIRL []. Recent work explores handling unstructured data [] and cross-embodiment scenarios []. Despite advancements, IRL faces challenges in computational efficiency and scalability []. RILe addresses these by jointly learning policy and reward function in a single process.

## 3 Background

### 3.1 Markov Decision Process

A standard Markov Decision Process (MDP) is defined by $(S,A,R,T,K,\gamma)$. $S$ is the state space that consists of all possible environment states $s$, and $A$ is action space that contains all possible environment actions $a$. $R=R(s,a):S\ \times\ A\ \rightarrow\ \mathbb{R}$ is the reward function. $T=\{P(\cdot|s,a)\}$ is the transition dynamics where $P(\cdot|s,a)$ is defined as an unknown state state transition probability function upon taking action $a\in A$ in state $s\in S$. $K(s)$ is the initial state distribution, i.e., $s_{0}\sim K(s)$ and $\gamma$ is the discount factor. The policy $\pi=\pi(a|s):S\ \rightarrow\ A$ is a mapping from states to actions. In this work, we consider $\gamma$-discounted infinite horizon settings. Following [], expectation with respect to the policy $\pi\in\Pi$ refers to the expectation when actions are sampled from $\pi(s)$: $E_{\pi}[R(s,a)]\triangleq E_{\pi}[\sum_{t=0}^{\infty}\gamma^{t}R(s_{t},a_{t})]$, where $s_{0}$ is sampled from an initial state distribution $\text{K}(s)$, $a_{t}$ is given by $\pi(\cdot|s_{t})$ and $s_{t+1}$ is determined by the unknown transition model as $P(\cdot|s_{t},a_{t})$. The unknown reward function $R(s,a)$ generates a reward given a state-action pair $(s,a)$. We consider a setting where $R=R(s,a)$ is parameterized by $\theta$ as $R_{\theta}(s,a)\in\mathbb{R}$ [].

Our work considers an imitation learning problem from expert trajectories, each consisting of states $s$ and actions $a$.
The set of expert trajectories $\tau_{E}$ are sampled from an expert policy $\pi_{E}\in\Pi$, where $\Pi$ is the set of all possible policies.
We assume that we have access to $m$ expert trajectories, all of which have $n$ time-steps, $\tau_{E}=\{(s^{i}_{0},a^{i}_{0}),(s^{i}_{1},a^{i}_{1}),\ldots,(s^{i}_{n},a^{i}
_{n})\}_{i=1}^{m}$.

### 3.2 Reinforcement Learning (RL)

Reinforcement learning seeks to find an optimal policy, $\pi^{*}$. that maximizes the discounted cumulative reward given from the reward function $R=R(s,a)$ (Fig. 1-(a)). In this work, we incorporate entropy regularization using the $\gamma$-discounted casual entropy function $H(\pi)=E_{\pi}[-\text{log}\ \pi(a|s)]$ [, ]. The RL problem with a parameterized reward function and entropy regularization is defined as

$$ $\text{RL}(R_{\theta}(s,a))=\pi^{*}=\operatorname*{arg\,max}_{\pi}E_{\pi}[R_{ \theta}(s,a)]+H(\pi).$ (1) $$

Figure: (a) RL
Refer to caption: extracted/5903656/algo_plots/rl_simple_wout_numbers.png

### 3.3 Inverse Reinforcement Learning (IRL)

Given sample trajectories $\tau_{E}$ from an optimal expert policy $\pi_{E}$, inverse reinforcement learning aims to recover a reward function $R^{*}_{\theta}(s,a)$ that maximally rewards the expert’s behavior (Fig. 1-(b)). Formally, IRL seeks a reward function, $R_{\theta}^{*}(s,a)$, satisfying: $E_{\pi_{E}}[\sum_{t=0}^{\infty}\gamma^{t}R_{\theta}^{*}(s_{t},a_{t})]\geq E_{
\pi}[\sum_{t=0}^{\infty}\gamma^{t}R_{\theta}^{*}(s_{t},a_{t})]\quad\forall\pi.$ Optimizing this reward function with reinforcement learning yields a policy that replicates expert behavior: $\text{RL}(R_{\theta}^{*}(s,a))=\pi^{*}=\operatorname*{arg\,max}_{\pi}E_{\pi}[
\sum_{t=0}^{\infty}\gamma^{t}R_{\theta}^{*}(s_{t},a_{t})]$.

Since only the expert’s trajectories are observed, expectations over $\pi_{E}$ are estimated from samples in $\tau_{E}$. Incorporating entropy regularization $H(\pi)$, maximum causal entropy inverse reinforcement learning [] is defined as

$$ $\text{IRL}(\tau_{E})=\operatorname*{arg\,max}_{R_{\theta}(s,a)\in\mathbb{R}} \left(E_{s,a\in\tau_{E}}[R_{\theta}(s,a)]-\max_{\pi}\left(E_{\pi}[R_{\theta}(s ,a)]+H(\pi)\right)\right).$ (2) $$

### 3.4 Adversarial Imitation Learning (AIL) and Adversarial Inverse Reinforcement Learning (AIRL)

Imitation Learning (IL) aims to directly approximate the expert policy from given expert trajectory samples $\tau_{E}$. It can be formulated as

$$ $\text{IL}(\tau_{E})=\operatorname*{arg\,min}_{\pi}E_{(s,a)\sim\tau_{E}}[L(\pi( \cdot|s),a)],$ (3) $$

where $L$ is a loss function, that captures the difference between policy and expert data.

GAIL [] introduces an adversarial imitation learning setting by quantifying the difference between the agent and the expert with a discriminator $D_{\phi}(s,a)$, parameterized by $\phi$ (Fig. 1-(c)).

The discriminator distinguishes between between expert-generated state-action pairs $(s,a)\sim\tau_{E}$ and non-expert state-action pairs $(s,a)\notin\tau_{E}$. The goal of GAIL is to find the optimal policy that fools the discriminator while maximizing an entropy constraint. The optimization is formulated as a zero-sum game between the discriminator $D_{\phi}(s,a)$ and the policy $\pi$:

$$ $\min_{\pi}\max_{\phi}E_{\pi}[\text{log}\ D_{\phi}(s,a)]+E_{\tau_{E}}[\text{log }\ (1-D_{\phi}(s,a))]-\lambda H(\pi).$ (4) $$

In other words, the reward function that is maximized by the policy is defined as a similarity function, expressed as $R(s,a)=-\text{log}\ (D_{\phi}(s,a))$.

AIRL [] extends this adversarial approach to inverse reinforcement learning, aiming to recover a reward function decoupled from environment dynamics (Fig. 1-(c)). The key insight of AIRL is to structure the discriminator as:

$$ $D_{\phi,\psi}(s,a,s^{\prime})=\frac{\exp(f_{\phi}(s,a,s^{\prime}))}{\exp(f_{ \phi}(s,a,s^{\prime}))+\pi(a|s)},$ (5) $$

where $f_{\phi}(s,a,s^{\prime})=r_{\psi}(s,a)+\gamma V_{\phi}(s^{\prime})-V_{\phi}(s)$. Here, $r_{\psi}(s,a)$ represents the learned reward function that is decoupled from the environment dynamics, $\gamma V_{\phi}(s^{\prime})-V_{\phi}(s)$. The AIRL optimization problem is formulated equivalently to GAIL (see Eqn. 4). The reward function $r_{\psi}(s,a)$ is learned through minimizing the cross-entropy loss inherent in this adversarial setup. Therefore, the reward function remains tightly coupled with the discriminator’s learning process.

## 4 RILe: Reinforced Imitation Learning

We propose Reinforced Imitation Learning (RILe) to learn the reward function and acquire a policy that emulates expert-like behavior simultaneously in one learning process.
Our framework consists of three key components: a student agent, a trainer agent, and a discriminator (Fig. 2).

Intuitively, the student agent learns the action policy and the trainer agent learns the reward function. Both agents are trained simultaneously via RL with the help of an adversarial discriminator.
Building on AIL, we utilize the discriminator to enable joint training of policy and reward function. Unlike traditional IRL, the discriminator provides immediate feedback to the trainer, eliminating the need to wait for the student agent’s convergence.

The trainer agent, trained via RL, is rewarded by the discriminator, and thus motivated to maximize the cumulative reward from the discriminator. The trainer learns to craft a reward function that promotes expert-like behavior in the student, even in complex scenarios. Crucially, the trainer’s RL-based approach optimizes for long-term performance rather than immediate feedback. This allows the trainer to encourage exploration and gradual improvement in three main ways. First, the trainer can associate its reward-giving actions with future improvements in the student’s behavior, even if those improvements take many steps to manifest. Second, the trainer can provide small positive rewards for actions that move the student in the right direction, even if they don’t immediately result in expert-like behavior. Finally, the trainer can adapt its reward function based on the student’s current capabilities, providing a learning path that gradually guides the student towards expert behavior. This approach enables RILe to overcome limitations of previous methods, particularly in complex tasks requiring extensive exploration, by promoting the discovery of expert-like strategies even when the student’s initial policy significantly diverges from expert behavior.

In the following, we define the components of RILe and explain how they can efficiently learn behavior from imperfect data.

Figure: Figure 2: Framework overview. 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.
Refer to caption: extracted/5903656/figures/RILe.png

#### Student Agent

The student agent aims to learn a policy $\pi_{S}$ by interacting with an environment in a standard RL setting within an MDP. For each of its actions $a^{S}\in A$, the environment returns a new state $s^{S}\in S$. However, rather than from a hand-crafted reward function, the student agent receives its reward from the policy of the trainer agent $\pi_{T}$. Therefore, the reward function is represented by the trainer policy. Thus, the student agent is guided by the actions of the trainer agent, i.e., the action of the trainer is the reward of the student: $r^{S}=\pi_{T}((s^{S},a^{S}))$. The optimization problem of the student agent is then defined as

$$ $\min_{\pi_{S}}-E_{(s^{S},a^{S})\sim\pi_{S}}[\pi_{T}\left((s^{S},a^{S})\right)].$ (6) $$

#### Discriminator

The discriminator differentiates between expert-generated state-action pairs $(s,a)\sim\tau_{E}$ and state-action pairs stemming from the policy of the student $(s,a)\sim\pi_{S}$.
In RILe, the discriminator is defined as a feed-forward deep neural network, parameterized by $\phi$.
Hence, the optimization problem is

$$ $\max_{\phi}E_{(s,a)\sim\tau_{E}}[\text{log}(D_{\phi}(s,a))]+E_{(s,a)\sim\pi_{S }}[\text{log}(1-D_{\phi}(s,a))].$ (7) $$

To provide effective guidance, the discriminator needs to accurately distinguish whether a given state-action pair originates from the expert distribution $(s,a)\sim\tau_{E}$ or not $(s,a)\notin\tau_{E}$.
The feasibility of this discrimination has been demonstrated by GAIL [].
The according lemma and proof are presented in the Appendix A.

#### Trainer Agent

The trainer agent aims to guide the student to approximate expert behavior by operating as its reward mechanism. Since the trainer agent takes the role of a reward function for the student, a new MDP is defined for the trainer agent: $\text{MDP}_{\text{T}}:(S_{T},A_{T},R_{T},T_{T},K,\gamma)$, where $S_{T}:S\ \text{x}\ A$ is the state space that consists all possible state action pairs from the standard MDP, defined in 3.1. $A_{T}$ is the action space, a mapping from $S_{T}\ \rightarrow\ \mathbb{R}$, so the action is a scalar value. $R_{T}(s,a)$ is the reward function where $s\in S_{T}$ and $a\in A_{T}$. $T_{T}=\{P(\cdot|s,a)\}$ is the transition dynamics where $P(\cdot|s,a)$ is defined as the state distribution upon taking action $a\in A_{T}$ in state $s\in S_{T}$. $K_{T}$ is the initial state distribution.

The trainer agent learns a policy $\pi_{T}$ that produces adequate reward signals to guide the student agent, by learning in a standard RL setting, within $MDP_{T}$. Since the state space of $MDP_{T}$ is defined over state-action pairs of $MDP_{S}$, the state of the trainer comprises the state-action pair of the student $s^{T}=(s^{S},a^{S})\in S_{T}$. It generates a scalar action $a^{T}$ which is given to the student agent as reward $r^{S}$, and bounded between -1 to 1. To help the trainer understand how its actions impact the reward it receives, we define the reward function such that it multiplies the scaled discriminator’s output by trainer’s actions. Therefore, the trainer agent’s reward function is defined as $R^{T}=\upsilon(D_{\phi}(s^{T}))(a^{T})$, where $D_{\phi}(s^{T})$ is the output of the discriminator and $\upsilon(x)=2x-1$ is the scaling function. By incorporating $a^{T}$ into the reward function, the trainer learns to adjust its policy based on the effectiveness of its previous actions.

The optimization problem of the trainer can be defined as

$$ $\max_{\pi_{T}}E_{\begin{subarray}{c}(s,a)\sim\pi_{S}\\ a^{T}\sim\pi_{T}\end{subarray}}[\upsilon(D_{\phi}(s,a))a^{T}].$ (8) $$

#### RILe

RILe combines the three components defined previously in order to find a student policy that mimics expert behaviors presented in $\tau_{E}$.
In RILe, the student policy $\pi_{S}$ and the trainer policy $\pi_{T}$ can be trained via any single-agent online reinforcement learning method. The training algorithm is given in Appendix D. Overall, the student agent aims to recover the optimal policy $\pi_{S}^{*}$ defined as

$$ $\pi_{S}^{*}=\operatorname*{arg\,max}_{\pi_{S}}E_{(s^{S},a^{S})\sim\pi_{S}} \left[\sum_{t=0}^{\infty}\gamma^{t}[\pi_{T}\left((s^{S}_{t},a^{S}_{t})\right) \right].$ (9) $$

At the same time, the trainer agent aims to recover the optimal policy $\pi_{T}^{*}$ as

$$ $\pi_{T}^{*}=\operatorname*{arg\,max}_{\pi_{T}}E_{\begin{subarray}{c}s^{T}\sim \pi_{S}\\ a^{T}\sim\pi_{T}\end{subarray}}\left[\sum_{t=0}^{\infty}\gamma^{t}[\upsilon(D_ {\phi}(s^{T}_{t}))a^{T}\right].$ (10) $$

To prove that the student agent can learn expert-like behavior, we need to show that the trainer agent learns to give higher rewards to student experiences that match with the expert state-action pair distribution, as this would enable a student policy to eventually mimic expert behavior.

Lemma 1: Given the discriminator $D_{\phi}$, the trainer agent optimizes its policy $\pi^{\theta_{T}}$ via policy gradients to provide rewards that guide the student agent to match expert’s state-action distributions.

Proof for Lemma 1
The assumptions are presented in Appendix A. The expert’s state-action distribution is denoted by $p_{\text{expert}}(s,a)$. The role of the trainer is to provide a reward signal to the student that encourages the approximation of $p_{\text{expert}}(s,a)$ as closely as possible.

We have $D_{\phi}:\mathcal{S}\times\mathcal{A}\rightarrow[0,1]$ as the discriminator, parameterized by $\phi$, which outputs the likelihood that a given state-action pair $(s,a)$ originates from the expert, as opposed to the student. The trainer’s policy $\pi^{\theta_{T}}$, parameterized by $\theta_{T}$, aims to maximize the likelihood under the discriminator’s assessment, thus encouraging the student agent to generate state-action pairs drawn from a distribution resembling $p_{\text{expert}}(s,a)$.

The Value and Q functions of the trainer, conditioned on the rewards provided by the discriminator, are defined in terms of expected cumulative discriminator rewards. The value function for the trainer’s policy parameters $\theta_{T}$ at state $s_{t}$, simplified from $s^{T}_{t}$ for better readability, is given by:

$$ $V^{\theta_{T}}(s_{t})=\mathbb{E}_{\pi^{\theta_{T}}}\left[\sum_{k=0}^{\infty} \gamma^{k}r_{t+k}\mid s_{t}\right],$ (11) $$

where $r_{t+k}=\upsilon(D_{\phi}(s_{t+k}))a_{t+k}$.
Similarly, the Q-function for taking action $a_{t}$ in state $s_{t}$ and then following policy $\pi^{\theta_{T}}$ can be written as:

$$ $Q^{\theta_{T}}(s_{t},a_{t})=D_{\phi}(s_{t})+\gamma\mathbb{E}_{\pi^{\theta_{T}} }\left[V^{\theta_{T}}(s_{t+1})\mid s_{t},a_{t}\right].$ (12) $$

The trainer’s policy optimization is done by maximizing the following clipped surrogate objective function:

$$ $\begin{split}L^{CLIP}(\theta_{T})=&\mathbb{E}_{(s_{t},a_{t})\sim\pi^{\theta_{T }}}\Big{[}\min\Big{(}\frac{\pi^{\theta_{T}}(a_{t}|s_{t})}{\pi^{\theta_{T_{old} }}(a_{t}|s_{t})}A^{\theta_{T_{old}}}(s_{t},a_{t}),\\ &\qquad\qquad\qquad\text{clip}(\frac{\pi^{\theta_{T}}(a_{t}|s_{t})}{\pi^{ \theta_{T_{old}}}(a_{t}|s_{t})},1-\epsilon,1+\epsilon)A^{\theta_{T_{old}}}(s_{ t},a_{t})\Big{)}\Big{]},\end{split}$ (13) $$

where $A^{\theta_{T_{old}}}(s_{t},a_{t})$ is the advantage function computed as $Q^{\theta_{T_{old}}}(s_{t},a_{t})-V^{\theta_{T_{old}}}(s_{t})$, defined with respect to the trainer’s old policy parameters $\theta_{T_{old}}$.

By expressing the advantage using the reward from the discriminator, we explicitly tie the policy gradient updates to the discriminator’s output, emphasizing the shaping of $\pi^{\theta_{S}}$ to match $p_{\text{expert}}(s,a)$:

$$ $A^{\theta_{T_{old}}}(s_{t},a_{t})=D_{\phi}(s_{t})+\gamma\mathbb{E}_{\pi^{ \theta_{S}}}\left[V^{\theta_{T_{old}}}(s_{t+1})\right]-V^{\theta_{T_{old}}}(s_ {t}).$ (14) $$

During each policy update, the objective in Equation 13 is maximized, driving parameter updates to favor actions that elicit higher rewards from the discriminator – effectively the actions that better align with expert behavior:

$$ $\theta_{T}\leftarrow\theta_{T}+\alpha_{\pi}\nabla_{\theta_{T}}L^{CLIP}(\theta_ {T}).$ (15) $$

The update rule in Equation 15, driven by cumulative rewards from the discriminator, incrementally adapts the trainer’s policy to reinforce student behaviors that are indistinguishable from those of the expert according to $D_{\phi}$. The trainer’s policy is guided to facilitate the student’s approximation of $p_{\text{expert}}(s,a)$, thereby fulfilling its role in the imitation learning process.

## 5 Experiments

We evaluate the performance of RILe by addressing three key questions:

- 1.
How does RILe’s adaptive reward function evolve compared to baseline methods and how does this evolution enhance the learning process?
- 2.
Is RILe efficient and scalable to high-dimensional continuous control tasks?
- 3.
Can RILe use expert-data explicitly to imitate expert behavior?

To answer the first question, we compare RILe’s performance with AI(R)L baselines in a maze setting, where we demonstrate how the trainer agent modifies the reward function to guide the student during training. For the second question, we evaluate RILe’s effectiveness in imitating motion-capture data within continuous robotic control tasks, using LocoMujoco [] as the environment. To answer the last question, we use a simulated humanoid character from MuJoCo [, ] to evaluate RILe’s performance when expert data is explicitly used by the agents. Additional experimental results are provided in the Appendix, where we evaluate the robustness of the learned reward function in changing environments and analyze the noise resilience of our method.

Baselines
We compare RILe with four baseline methods: Behavioral cloning (BC [, ]), adversarial imitation learning (GAIL []), adversarial inverse reinforcement learning (AIRL []) and inverse reinforcement learning (IQ-Learn []).

### 5.1 Evolving Reward Function

To evaluate the impact of RILe’s trainer agent on the learning process compared to existing methods, we designed a maze experiment. The task is to navigate from a fixed start on the left side to a goal in the top-right corner, avoiding obstacles.

Using a single expert demonstration, we trained RILe, GAIL, and AIRL. Fig. 3 shows how the student’s reward function evolves during training with all three methods. For RILe, we plot the reward function learned by the trainer. For GAIL, we visualize the discriminator output, and for AIRL, the reward term under the discriminator. The columns represent reward plots at 25%, 50%, 75%, and 100% of training completion for each algorithm, respectively. We present the student agent’s trajectory from the previous epoch on reward plots to demonstrate how reward functions adapt to the student agent’s progress.

Figure: (a) RILe
Refer to caption: extracted/5903656/maze_plots/rile_all.jpg

The discriminator-based reward functions of GAIL and AIRL remain relatively static, largely independent of the student agent’s current performance. This static behavior is because the discriminators in GAIL and AIRL focus on distinguishing expert data and progressively improve throughout training. In contrast, the trainer in RILe adapts itself based on the current learning progress of the student’s policy.

As the first column shows, when the student’s trajectory does not resemble the expert, RILe’s trainer encourage exploration towards the expert path. This indicates that the trainer can provide informative rewards despite potentially low discriminator feedback. Then, in the second column, as the student progresses, RILe’s reward function adapts dynamically. When the student learns to reach the bottom-right, the trainer shifts high rewards to the top-left, guiding the agent to explore that area. Third column shows as the student approaches the goal, the trainer increases rewards around it while maintaining rewards in specific areas (e.g., the left part) to prevent the agent from getting stuck. Finally, as the fourth column shows, the trainer focuses high rewards on the top-right area near the goal while maintaining a path of rewards that guides the student along the expert trajectory.

RILe’s evolving reward function demonstrates its ability to provide meaningful guidance even when the discriminator easily identifies non-expert policies. By adapting to the student’s current capabilities, RILe creates a dynamic learning curriculum that encourages exploration and gradual improvement towards expert-like behavior. Our results demonstrate that the trainer effectively associates rewards with future behavior improvements, even if delayed. It provides incremental rewards for actions moving the student in the right direction, and adapts the reward function to match the student’s abilities, gradually guiding them toward expert behavior.

### 5.2 Motion-Capture Data Imitation for Robotic Continuous Control

We evaluate RILe’s performance on the LocoMujoco benchmark, which presents a challenging task of imitating motion-capture (MoCap) data for various robotic control tasks. This benchmark is particularly demanding due to its high dimensionality and the absence of action data in the MoCap recordings which precludes the use of methods like Behavioral Cloning and IQLearn that require complete state-action pairs.

**Table 1: Test and validation results on seven LocoMujoco tasks.**
|  | Test Seeds | Validation Seeds |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
|  | RILe | GAIL | AIRL | RILe | GAIL | AIRL | Expert |
| Atlas Walk | 870.6 | 792.7 | 300.5 | 895.4 | 918.6 | 356 | 1000 |
| Atlas Carry | 850.8 | 669.3 | 256.4 | 889.7 | 974.2 | 271.9 | 1000 |
| Talos Walk | 842.5 | 442.3 | 102.1 | 884.7 | 675.5 | 103.4 | 1000 |
| Talos Carry | 220.1 | 186.3 | 134.2 | 503.3 | 338.5 | 74.1 | 1000 |
| UnitreeH1 Walk | 966.2 | 950.2 | 568.1 | 980.7 | 965.1 | 716.2 | 1000 |
| UnitreeH1 Carry | 788.3 | 634.6 | 130.5 | 850.6 | 637.4 | 140.9 | 1000 |
| Humanoid Walk | 831.3 | 181.4 | 80.1 | 970.3 | 216.2 | 78.2 | 1000 |

Table 1 presents the results for seven LocoMujoco tasks, comparing RILe with GAIL and AIRL across both test and validation seeds. The expert performance, set at 1000 for all tasks, serves as the upper benchmark. RILe demonstrates superior performance in all scenarios, particularly excelling in generalization to new initial conditions as evidenced by the test seed results.

### 5.3 Impact of Expert Data on Trainer-Student Dynamics

Figure: Figure 4: Impact of explicit usage of expert data on RILe’s performance. Red markers and yellow markers show normalized scores and normalized steps, respectively. Square markers show IQ-Learn and BC results for comparisons. The percentage of expert data in the buffers are presented around the figure. More expert data in buffers speeds convergence but can reduce final performance, especially at 100%, showing how the explicit usage of excess expert data hinders trainer-student adaptation in RILe
Refer to caption: x1.png

We study how explicitly incorporating expert data into RILe’s training affects the trainer’s ability to adapt to the student’s needs, analyzing performance and convergence time when expert data is added to the agents’ replay buffers. For this experiment, we used MuJoCo’s Humanoid environment [, ] and a single expert trajectory from []. We varied the proportion of expert data in the replay buffers from 0% to 100%; for example, 25% means a quarter of the buffer is expert data and 75% is from agent interactions.

As shown in Fig. 4, introducing expert data led to faster convergence but decreased performance. Notably, when environmental interactions were completely replaced by the expert data (100% case), the student’s performance declined significantly. Excessive expert data appears to hinder the trainer’s ability to adapt to the student, disrupting RILe’s dynamic learning process. For comparison, we include results from IQ-Learn and Behavioral Cloning (BC), which rely explicitly on expert data. Despite their exclusive use of expert data, neither method matched RILe’s performance, even when RILe used substantial amounts of expert data.

Our findings underscore the importance of balancing environmental interactions and expert demonstrations in buffers in RILe. While some expert data accelerates learning, too much impedes the trainer’s adaptability, underscoring the unique strength of RILe’s interactive learning approach.

## 6 Discussion

Our experiments have demonstrated that RILe consistently outperforms baseline models across various settings, surpassing state-of-the-art adversarial IL/IRL methods. RILe’s key strength lies in its adaptive learning approach, where the trainer agent dynamically adjusts the reward function based on the student’s current capabilities. This enables more efficient learning, particularly in complex environments.

However, RILe faces challenges, primarily in maintaining policy stability with a changing reward function. Our solution of freezing the trainer after the student achieves half the expert score is effective but potentially limiting. Additionally, the discriminator’s tendency to overfit quickly presents another significant challenge. Future work should focus on establishing bounds for reward agent updates to enhance stability without compromising adaptability, developing a stronger theoretical understanding of RILe’s convergence properties, and exploring discriminator-less approaches.

Despite these challenges, RILe has demonstrated advantages in adaptability, robustness, and generalization, thereby making it a promising direction for future research in imitation learning and opening up new possibilities for dynamic and responsive learning processes.

## Appendix A Justification of RILe

Assumptions:

- •
The discriminator loss curve is complex and the discriminator function, $D_{\phi}(s,a)$, is sufficiently expressive since it is parameterized by a neural network with adequate capacity.
- •
For the trainer’s and student’s policy functions $(\pi^{\theta_{T}})$ and $(\pi^{\theta_{S}})$, and the Q-functions $(Q^{\theta_{S}})$, each is Lipschitz continuous with respect to its parameters with constants $(L_{\theta_{T}}),(L_{\theta_{S}}),and(L_{Q})$, respectively. This means for all $(s,a)$ and for any pair of parameter settings $(\theta,\theta^{\prime}):[|\pi^{\theta}(s,a)-\pi^{\theta^{\prime}}(s,a)|\leq L
_{\theta}|\theta-\theta^{\prime}|,][|Q^{\theta}(s,a)-Q^{\theta^{\prime}}(s,a)|
\leq L_{Q}|\theta-\theta^{\prime}|.]$

### A.1 Lemma 2:

The discriminator $D_{\phi}$, parameterized by $\phi$ will converge to a function that estimates the probability of a state-action pair being generated by the expert policy, when trained on samples generated by both a student policy $\pi^{\theta_{S}}$ and an expert policy $\pi_{E}$.

Proof for Lemma 2:
The training objective for the discriminator is framed as a binary classification problem over mini-batches from expert demonstrations and student-generated trajectories. The discriminator’s loss function $\mathcal{L}_{D}(\phi)$ is the binary cross-entropy loss, which for a mini-batch of size $n$ is defined as:

$$ $\mathcal{L}_{D}(\phi)=-\frac{1}{n}\sum_{i=1}^{n}\Big{[}y_{i}\log(D_{\phi}(s_{i },a_{i}))+(1-y_{i})\log(1-D_{\phi}(s_{i},a_{i}))\Big{]},$ (16) $$

where $(s_{i},a_{i})$ are sampled state-action pairs from the combined replay buffer $\mathcal{D}=\mathcal{D}_{S}\cup\mathcal{D}_{E}$, with corresponding labels $y_{i}$ indicating whether the pair is from the expert $(y_{i}=1)$ or the student $(y_{i}=0)$.
The stochastic gradient descent update rule for optimizing $\mathcal{L}_{D}(\phi)$ is then given by:

$$ $\phi\leftarrow\phi-\eta_{D}\nabla_{\phi}\mathcal{L}_{D}(\phi),$ (17) $$

where $\eta_{D}$ is the learning rate for the discriminator.

Through iterative updates, $D_{\phi}(s,a)$ will converge to $P(\pi_{E}|s,a)$, provided the minimization of $\mathcal{L}_{D}(\phi)$ progresses according to the theoretical foundations of stochastic gradient descent. The convergence relies on the assumption that the student’s policy $\pi^{\theta_{S}}$ and the expert policy $\pi_{E}$ induce stationary distributions of state-action pairs, such that the discriminator’s data source is consistently representing both policies over time.

By minimizing $\mathcal{L}_{D}(\phi)$, we seek $\phi^{*}$ such that:

$$ $\phi^{*}=\arg\min_{\phi}\mathcal{L}_{D}(\phi).$ (18) $$

Under typical conditions for convergence in stochastic gradient descent, the convergence to a local minimum or saddle point can be guaranteed. The discriminator’s ability to distinguish between student and expert pairs improves as $\mathcal{L}_{D}(\phi)$ is minimized, implying that $\lim_{n_{batch}\rightarrow\infty}D{\phi^{*}}(s,a)=P(\pi_{E}|s,a)$, where $n_{batch}$ is the number of batches.

## Appendix B Additional Experiments

### B.1 Noisy Expert Data

Figure: (a) Expert traj.
Refer to caption: extracted/5903656/figures/expert_grid_5_noborder.png

To demonstrate the advantage of using RL to learn the reward function in RILe, as opposed to deriving the reward directly from the discriminator in AIL and AIRL, we designed a 5x5 MiniGrid experiment.
The grid consists of 4 lava tiles that immediately kill the agent if it steps in it, representing terminal conditions.
The goal condition of the environment is reaching the green tile.

The expert demonstrations are imperfect, depicting an expert that passes through a lava tile without being killed and still reaches the green goal tile.
Using this data, we trained the adversarial approaches with a perfect discriminator, which provides a reward of $0.99$ if the visited state-action pair stems from the expert and $0.01$ otherwise.
These values were chosen over 1 and 0 because both AIRL and GAIL use the logarithm of the discriminator output to calculate rewards.

Results are presented in Fig. 5. The value graphs (Fig. 5e-g) are attained by computing the value of each grid cell $c_{i}$ as $\sum_{a\in A}D(c_{i},a)$ for AIRL and GAIL, and $\sum_{a\in A}\pi_{T}(c_{i},a)$ for RILe. Fig. 5a shows the expert trajectory.

GAIL (Fig. 5c), AIRL (Fig. 5d) and IQLearn (Fig. 5e) fail to reach the goal, as their agents either become stuck or are directed into lava.

In contrast, RILe (Fig. 5d) successfully reaches the goal, demonstrating its ability to navigate around imperfections in expert data.
The difference in the value graphs between RILe and the baselines intuitively explains this outcome.
In AIL and AIRL (Fig. 5f-g), the optimal paths, defined by the actions most rewarded by their discriminators, follow the noisy expert data perfectly.
Similarly, in IQLearn, the agent tries to match expert state-actions as closely as possible, minimizing any deviation from the expert trajectory.
In contrast, RILe’s trainer agent, trained using RL, adds an extra degree of freedom in the adversarial IL/IRL setting.
By providing rewards that maximize cumulative returns from the discriminator, rather than deriving the reward directly from its output, the value graph (Fig. 5f) can learn to circumvent the lava tile in order to follow the expert trajectory to the goal.
Consequently, the optimal path of the student agent can overcome the sub-optimal state suggested by the noisy expert demonstration. Since the student agent is guided by the trainer to also match the expert trajectory, it remains close to this path after passing the lava tiles.

### B.2 Robustness to Noise in the Expert Data

To evaluate the robustness of RILe and baseline methods to noise in the expert data, we conducted experiments in the MuJoCo Humanoid-v2 environment. Artificial noise sampled from a zero-mean Gaussian distribution with varying standard deviations ($\Sigma$) was added to a single expert trajectory, affecting either the actions or the states. The baselines used for comparison were GAIL [], AIRL [], RIL-Co [], IC-GAIL [], and IQ-Learn [].

As shown in Table 2, RILe consistently outperforms the baselines across different noise levels, demonstrating superior robustness even when a high amount of noise is present in the expert data ($\Sigma=0.5$). These results indicate that RILe is less sensitive to imperfections in the expert demonstrations compared to existing methods.

**Table 2: Test results in MuJoCo Humanoid-v2 environment, where artificial noise sampled from a zero-mean Gaussian distribution is added to a single expert trajectory. Results are aggregated over 20 different-seed environments. IQ-Learn* is trained using the official code and hyperparameters of the IQ-Learn algorithm.**
|  | Noise-Free | Action Noise | State Noise |  |  |
| --- | --- | --- | --- | --- | --- |
|  | $\Sigma=0$ | $\Sigma=0.2$ | $\Sigma=0.5$ | $\Sigma=0.2$ | $\Sigma=0.5$ |
| RILe | 5681 | 5280 | 5154 | 5350 | 5205 |
| GAIL | 5430 | 5275 | 902 | 5147 | 917 |
| AIRL | 5276 | 4869 | 4589 | 4898 | 4780 |
| RIL-Co | 576 | 491 | 493 | 505 | 501 |
| IC-GAIL | 610 | 601 | 568 | 590 | 591 |
| IQ-Learn* | 312 | 192 | 153 | 243 | 277 |

### B.3 Robustness of the Learned Reward Function

We evaluated the robustness of the reward functions learned by RILe and AIRL [] through an experiment similar to that conducted by []. Initially, both methods were trained to learn reward functions in a noise-free MuJoCo Humanoid-v2 environment. After training, these reward functions were frozen. Subsequently, new student agents were trained using these fixed reward functions in environments where Gaussian noise was added to the agents’ actions, with varying noise levels.

Table 3 presents the results of this evaluation. The reward function learned by RILe demonstrates superior robustness to noise, maintaining high performance even under increased noise levels. In contrast, the performance of agents using the reward function learned by AIRL decreases more significantly as noise increases. These findings indicate that the reward function learned by RILe is more resilient to environmental noise, contributing to better agent performance in noisy conditions.

**Table 3: We test the robustness of learned reward functions. After training reward functions in a noise-free setting, reward functions are frozen, and used to train a new agent in a noisy environment, where Gaussian noise is added to agent’s actions in every step.**
|  | No Noise | Mild Noise | High Noise |
| --- | --- | --- | --- |
|  | $\Sigma=0$ | $\Sigma=0.2$ | $\Sigma=0.5$ |
| RILe | 5748 | 5201 | 5196 |
| AIRL | 5334 | 5005 | 4967 |

### B.4 Reward Curves

We compare the reward curves of RILe, GAIL [], AIRL [], IQ-Learn [], and AdapMen [] in the MuJoCo Humanoid-v2 experiment. Since the task involves learning from expert trajectories, we combined AdapMen with an adversarial discriminator to enable training without an expert policy.

As shown in the reward curves, despite RILe having multiple components, it is the most efficient method. This efficiency is achieved through the dynamic guidance of the trainer during training, which adapts the reward function to meet the student’s needs.

Figure: Figure 6: Training reward curves for the MuJoCo Humanoid-v2 experiment comparing RILe, AIRL, GAIL, IQ-Learn*, and adapted AdapMen. AdapMen is combined with an adversarial discriminator to be able to train it without expert policy.
Refer to caption: extracted/5903656/figures/reward_curve.png

## Appendix C Compute Resources

For the training of RILe and baselines, following computational sources are employed:

- •
AMD EPYC 7742 64-Core Processor
- •
1 x Nvidia A100 GPU
- •
32GB Memory

## Appendix D Algorithm

Figure: Algorithm 1 RILe Training Process

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. The reviewer's primary concern is how the trainer's reward function contributes to optimizing long-term feedback. Additionally, the reviewer questions the rationale for incorporating the trainer's action into the reward function. Please refer to the "Weaknesses" section for a detailed explanation.  
2. L313, how is the constant determined in $v(x) = 2x - 1$?  
3. How many training environment steps are used for each task?   
4. Why are confidence intervals not included in Table 1?
soundness: 3
strengths: **Originality & Significance** 

1. The reviewer found the concepts of a custom reward function and the decoupling of reward function learning from the student and discriminator to be quite innovative. The experimental results suggest that this approach effectively addresses the issue of insufficient guidance, a common challenge in early-stage training for adversarial imitation and reinforcement learning methods. This improvement in guidance appears to be particularly beneficial during the initial phases of the learning process.

2. The background section is well-organized and effectively introduces the key concepts. The authors provide a concise overview of reinforcement learning (RL), inverse reinforcement learning (IRL), adversarial imitation learning (AIL), and adversarial reinforcement learning (AIRL). These explanations establish a solid foundation for understanding the proposed approach and its position within the broader literature. Additionally, the authors clearly articulate the distinctions between existing methods and their proposed approach, further highlighting the motivation and novelty of their work. 


**Quality** 

The claims are supported by experimental results. Particularly, table 1 shows that the proposed approach is able to handle complex tasks that baselines didn’t perform well. 
The reviewer appreciated the presentation of figure 3. It clearly illustrates the dynamics of the reward function of the proposed approach and its difference to existing methods. 

**Clarity** 

While some clarification of  the technical contents is needed,  the paper is overall well organized and easy to follow. Please see the weakness section for details.
summary: This paper introduces ‘Reinforced Imitation Learning’, a new trainer-student framework. In this framework, the trainer receives feedback from a discriminator and continuously adjusts the output reward in response to the student's behavior. The student is then trained using reinforcement learning based on the reward signal from the trainer. This approach may enable increased exploration behavior in students and more detailed feedback from the trainer. The proposed approach is evaluated in a maze setting, LocoMujoco, and humanoid MuJoCo. Experimental results indicate that the proposed approach outperforms competitive baselines significantly.
weaknesses: 1. The reviewer found the reward function for the trainer agent needs further clarification. Specifically, L 241 reads ‘the trainer’s RL-based approach optimizes for long-term performance rather than immediate feedback’, but it is unclear to the reader how the formulation of the reward function (Eq. 8) could achieve it. 

2. It is unclear to the reviewer the role of $a^T$ in the reward function of the trainer (Eq. 8). With $a^T$ in the objective, would it encourage the trainer to generate $a^T=1$ without further constraint? Also, L 314 reads ‘by incorporating $a^T$ into the reward function, the trainer learns to adjust its policy based on the effectiveness of its previous actions’. It is unclear to the reviewer how this is achieved. Please elaborate.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. What is the rationale behind the specific definition of the optimization objective of the trainer agent, $v(D_{\phi}(s, a)) a^T$ (Equation 8)? It is mentioned that "By incorporating $a^T$ into the reward function, the trainer learns to adjust its policy based on the effectiveness of its previous actions", but this explanation is somewhat unclear. Could you clarify why $a^T$ is specifically multiplied with $v(D_{\phi}(s, a))$?

2. In Figure 4, the performance of RILe remains relatively good when the trainer's replay buffer is 100\% expert data, provided that the proportion of expert data in the student's replay buffer is not too large. This outcome is somewhat puzzling, as I would expect that with 100\% expert data in the trainer's replay buffer, neither the trainer nor the discriminator would be trained on any trajectories generated by the student policy, based on the framework illustrated in Figure 2. Am I misunderstanding something here?
soundness: 1
strengths: 1. This paper aims to address some of the challenges faced by existing imitation learning methods in complex settings, a valuable problem to study with a wide range of potential practical applications.

2. The introduction of an additional trainer agent in the proposed method, RILe, is a novel design.

3. The discussion of related literature is clear and effectively situates the proposed method within the context of existing research.
summary: This paper proposes a novel method for adversarial imitation learning. Unlike established methods such as GAIL and AIRL, which learn a policy using reinforcement learning (RL) under a reward function derived directly from the discriminator, the proposed method, RILe, introduces an additional RL component called the trainer agent. This trainer agent generates the reward function for policy learning, effectively serving as an intermediary layer between the discriminator and the policy learning component (referred to as the student agent). The paper claims that this additional trainer component, as an RL agent, is able to dynamically provide an adaptive reward function throughout the student agent's policy learning process, resulting in superior imitation learning performance in complex settings.
weaknesses: 1. My major concern with the paper is that several major claims regarding the capabilities of the proposed method, RILe, are not sufficiently supported by rigorous analysis or adequate experimental evidence.
* Claim of RILe's adaptability: For instance, in Section 5.1, it is stated that "RILe demonstrates a more adaptive reward function that evolves with the student's progress ...". However, this claim is inadequately supported by the limited results in Figure 3, which only includes a single, simple setting and is open to subjective interpretation. Similar claims in the Introduction and Discussion suggest that RILe’s adaptability leads to more efficient policy learning for the student agent, but there is insufficient experimental evidence to validate this connection.
* Claim of RILe outperforming state-of-the-art adversarial IL methods: Although the main results in Section 5.2 show that RILe outperforms the baseline methods, GAIL and AIRL, on several LocoMujuco tasks, these results lack sufficient details on the experimental setup, such as the settings of tasks used for evaluation and the hyperparameter values of the compared methods.
* Additionally, the proof of the theoretical argument, Lemma 1, is neither rigorously presented nor complete.

2. Another concern I have with the proposed RILe method is that it stacks one RL agent (the trainer agent) on top of another (the student agent). The trainer agent's state transition depends on the student agent's policy since its states are represented by the state-action pairs generated by the student agent. Meanwhile, the student agent's policy is influenced by the trainer agent's policy through the reward function generation. This structure creates a dependency in the trainer's RL environment (i.e. the state transition of the RL environment faced by the trainer agent is not exogonous to its policy), making the MDP ill-defined and potentially introducing extra instability to the training framework.
* In fact, the paper briefly addresses training instability in the Discussion section and mentions a workaround of freezing the trainer midway through training.

3. The proposed optimization objective of the trainer agent, Equation 8, lacks a clear explanation of the rationale behind this particular choice. (See Question 1 below)

4. Overall, the paper is not well-organized and includes errors or inaccuracies in some parts.
* In Section 3.2, entropy regularization is introduced as part of the RL problem formulation in this work, defined in Equation 1. However, in Section 3.3, the RL formulation is presented without entropy regularization, which is then added again in Equation 2.
* In Section 4, the first term on the right-hand side of Equation 14 should be $v(D_{\phi}(s_t))a_t$ instead of $D_{\phi}(s_t)$ to remain consistent with the reward definition in Equation 11.
* The workaround for the training instability (freezing the trainer after meeting certain criteria) is only mentioned in the Discussion section, with no mention of it in the discussion of experimental settings and results.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. The experimental setup relies primarily on older baselines (e.g., BC from 2010, GAIL from 2016, and AIRL from 2018). Including comparisons with more recent adversarial imitation learning (AIL) methods, such as DRAIL [1], would provide a more comprehensive assessment of RILe’s performance, especially given that DRAIL also address challenges related to learning efficiency in complex, high-dimensional environments. The paper needs to include DRAIL [1]. Doing so would significantly enhance the paper’s impact. I would be willing to score it higher with this addition.

2. How does the dynamic reward adjustment by the trainer agent handle rapidly changing or highly stochastic environments, where optimal behavior may vary significantly within episodes?

3. RILe involves multiple components with different objectives (trainer, student, and discriminator). How sensitive is RILe to hyperparameter choices across these agents, and are there any recommended tuning strategies?

4.  The paper shows RILe’s ability to perform well with noisy expert data, but how robust is this approach if the noise distribution changes dynamically during training, or if the expert data is sparse?

[1] Lai, Chun-Mao, et al. "Diffusion-Reward Adversarial Imitation Learning." arXiv preprint arXiv:2405.16194 (2024).
soundness: 3
strengths: 1. **Motivation and Intuition**: The motivation for reducing reliance on static discriminators is compelling, as complex tasks often necessitate adaptive guidance.

2. **Novelty**: The framework’s use of a trainer agent that dynamically learns a reward function to align with student policy is an innovative approach that enhances traditional adversarial learning strategies.

3. **Technical Contribution**: RILe’s decoupling of the reward function learning from the student’s policy learning shows clear benefit, especially in tasks requiring extensive exploration, as the framework enables more progressive feedback than binary classifications.

4. **Clarity**: The paper is well-structured, with clear theoretical explanations and effective visualizations. Figures, such as Figure 3, successfully illustrate the evolution of reward functions across training stages.
summary: RILe (Reinforced Imitation Learning environment) is a framework that enhances reinforcement and imitation learning by introducing a dynamic, adaptive reward function responsive to agent performance. Addressing limitations of traditional IL methods like GAIL and AIRL, RILe’s trainer-student model continuously optimizes rewards through reinforcement learning, effectively guiding the student agent to mimic expert demonstrations in complex tasks. Experiments in maze and humanoid locomotion tasks demonstrate RILe's advantages in scalability, robustness, and adaptability, highlighting its effectiveness for high-dimensional control environments.
weaknesses: 1. **Limited Baselines** (important): The experiment lacks more recent baselines (BC from 2010, GAIL from 2016, and AIRL from 2018). It would strengthen the work to include comparisons with more recent AIL methods, such as DRAIL [1], V-MAIL [2], or ARC [3]. Especially, they should cite DRAIL [1], which also focus on improving learning efficiency, particularly in complex, high-dimensional environments like RILe does.

2. **Lack of Ablation Study**: The paper would benefit from additional ablation studies to better analyze the contributions of individual components within the RILe framework.

3. **Limited Task Diversity**: RILe is tested only on LocoMujoco (locomotion tasks). Expanding experiments to include navigation or manipulation tasks would improve the generalizability of the framework.

[1] Lai, Chun-Mao, et al. "Diffusion-Reward Adversarial Imitation Learning." arXiv preprint arXiv:2405.16194 (2024).

[2] Rafailov, Rafael, et al. "Visual adversarial imitation learning using variational models." Advances in Neural Information Processing Systems 34 (2021): 3016-3028.

[3] Deka, Ankur, Changliu Liu, and Katia P. Sycara. "ARC-Actor Residual Critic for Adversarial Imitation Learning." Conference on Robot Learning. PMLR, 2023.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: - In line 241, the authors mention "the trainer’s RL-based approach optimizes for long-term performance rather than immediate feedback".  If the benefit comes from RL, all AIL methods also share it since they all use RL for policy learning. How does the RILe handle this?

- In Section 5.1 and Figure 3, it’s stated that RILe offers better rewards for guiding student policies, but in the second and third columns, GAIL appears to learn faster. The reward contours for GAIL and AIRL also seem closer to the ground truth. Moreover, learning with a fixed reward function is the most common approach in typical RL. Could the authors elaborate more on how the evolving feature benefits policy learning?

- Why did the authors choose LocoMujoco, which lacks action labels, rather than more widely used Mujoco tasks for Learning from Demonstration (LfD)? I would recommend limiting the scope to Learning from Demonstration (LfD), where action labels are available. Besides, IQLearn did provide learning with state-only rewards version in the paper. If the authors want to demonstrate the application on Learning from Observation (LfO), I think it will be better to conduct experiments separately with a different set of baselines [5,6,7,8].

- In Figure 4, how are normalized scores and steps defined? The result seems to be applicable to many online imitation learning methods. Could the authors provide more details on the purpose of these metrics in the context of RILe?

- The authors mentioned RILe improves the computational efficiency compared to other IRL works. However, RILe needs to train an extra trainer network compared to GAIL. Is there a quantitative analysis on this?

References:

[1] Pomerleau, D. A. (1991). Efficient training of artificial neural networks for autonomous navigation. Neural computation, 3(1), 88-97.

[2] Papagiannis, G., & Li, Y. (2022, September). Imitation learning with sinkhorn distances. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases (pp. 116-131). Cham: Springer Nature Switzerland.

[3] Chi, C., Xu, Z., Feng, S., Cousineau, E., Du, Y., Burchfiel, B., ... & Song, S. (2023). Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal of Robotics Research, 02783649241273668.

[4] Lai, C. M., Wang, H. C., Hsieh, P. C., Wang, Y. C. F., Chen, M. H., & Sun, S. H. (2024). Diffusion-Reward Adversarial Imitation Learning. arXiv preprint arXiv:2405.16194.

[5] Torabi, F., Warnell, G., & Stone, P. (2018). Behavioral cloning from observation. arXiv preprint arXiv:1805.01954.

[6] Liu, M., Zhu, Z., Zhuang, Y., Zhang, W., Hao, J., Yu, Y., & Wang, J. (2022). Plan your target and learn your skills: Transferable state-only imitation learning via decoupled policy optimization. arXiv preprint arXiv:2203.02214.

[7] Torabi, F., Warnell, G., & Stone, P. (2018). Generative adversarial imitation from observation. arXiv preprint arXiv:1807.06158.

[8] Huang, B. R., Yang, C. K., Lai, C. M., Wu, D. J., & Sun, S. H. (2024). Diffusion Imitation from Observation. arXiv preprint arXiv:2410.05429.
soundness: 2
strengths: - The idea of Decoupled Reward-function Learning with AIL is novel.

- RILe demonstrates better performance on the LocoMujoco benchmark than baselines (GAIL and AIRL).
summary: The paper proposed Reinforced Imitation Learning (RILe) to learn the reward function for policy learning with AIL and teacher-student framework. The method is evaluated on LocoMujoco benchmark and it outperforms GAIL and AIRL.
weaknesses: - The connection between motivation and the proposed method is weak. More intuition or theoretical proofs are needed.
- While the paper highlights some issues with existing methods, it lacks detailed references and experiments to convincingly demonstrate how RILe addresses these issues. More ablation studies or theoretical analysis could strengthen the argument.
- The scope of the experiment is limited to comparing only GAIL and AIRL on one benchmark. More baselines [1,2,3,4] and tasks should be evaluated.

Meta-review: This paper proposes a novel method for adversarial imitation learning with three components: a discriminator that provides feedback to a trainer (an RL agent) which is optimized to reward the student (another RL agent). The paper argues that the additional trainer component can adapt the reward function to accelerate the student agent's policy learning process, perhaps increasing exploration. Experiments demonstrate strong results in locomotion tasks.

Reviewers appreciated the novelty of the proposed approach and strong empirical performance (esp. during the early stages of learning).
Authors appreciated the clear discussion of the relationship with prior work and overall clear writing and motivation.

Several reviewers recommended additional baselines, ablation experiments, and evaluation tasks (e.g., manipulation). They also provided feedback on paper organization and suggested that the paper might be strengthened by including additional theoretical results. Reviewers also had questions about why the reward function (Eq. 8) resulted in maximizing for long term performance and details of the experimental setup. Reviewers also suggested that some of the claims should be clarified (e.g., is it the _adaptability_ that is driving improved performance?) and had concerns about the complexity of the proposed approach, as it seems to stack on RL agent on top of another. (this is briefly discussed in the paper's Discussion).

Overall, despite the novelty of the proposed approach, concerns about the experiments (i.e., whether they support the specific claims made in the paper) compel me to recommend that this paper not be accepted.

Decision: Reject