On Discovering Algorithms for Adversarial Imitation Learning
Authors: Shashank Reddy Chirra
Jayden Teoh
Praveen Paruchuri
Pradeep Varakantham
Venue: ICLR 2026 Poster (2026)
Keywords: imitation learning
algorithm discovery
llms
evolutionary algorithms
Primary area: reinforcement learning
TL;DR: We use LLM-guided evolutionary search to automatically discover reward assignment functions that improve the stability and performance of adversarial imitation learning.

Abstract: Adversarial Imitation Learning (AIL) methods, while effective in settings with limited expert demonstrations, are often considered unstable. These approaches typically decompose into two components: Density Ratio (DR) estimation $\frac{\rho_E}{\rho_{\pi}}$, where a discriminator estimates the relative occupancy of state-action pairs under the policy versus the expert; and Reward Assignment (RA), where this ratio is transformed into a reward signal used to train the policy. While significant research has focused on improving density estimation, the role of reward assignment in influencing training dynamics and final policy performance has been largely overlooked. RA functions in AIL are typically derived from divergence minimization objectives, relying heavily on human design and ingenuity. In this work, we take a different approach: we investigate the discovery of data-driven RA functions, i.e, based directly on the performance of the resulting imitation policy. To this end, we leverage an LLM-guided evolutionary framework that efficiently explores the space of RA functions, yielding _Discovered Adversarial Imitation Learning_ (DAIL), the first meta-learnt AIL algorithm. Remarkably, DAIL generalises across unseen environments and policy optimization algorithms, outperforming the current state-of-the-art of _human-designed_  baselines. Finally, we analyse why DAIL leads to more stable training, offering novel insights into the role of RA functions in the stability of AIL.

Published as a conference paper at ICLR 2026
ON DISCOVERING ALGORITHMS FOR ADVERSARIAL IMITATION LEARNING
Shashank Reddy Chirra
Jayden Teoh
University of Oxford shashank@ rowots...tt....oooo ac.uk
Singapore Management University jxteoh. 20230 wmw.smmus........ edu. sg
Praveen Paruchuri
Pradeep Varakantham
III, Hyderabad
Singapore Management University pradeepve smu, edu. sq
praveeneilit.
ABSTRACT
Adversarial Imitation Learning (AIL) methods, while effective settings with limited expert demonstrations, are often considered unstable. These approaches typically decompose into two components: Density Ratio (DR) estimation es A. where a discriminator estimates the relative occupancy of state-action pairs un- der the policy versus the expert; and Reward Assignment (RA), where this ra- tio is transformed into a reward signal used to train the policy. While signifi- cant research has focused on improving density estimation, the role of reward as- signment in influencing training dynamics and final policy performance has been largely overlooked. RA functions in AIL are typically derived from divergence minimization objectives, relying heavily on human design and ingenuity. In this work, we take different approach: we investigate the discovery of data-driven RA functions, i.e, based directly on the performance of the resulting imitation policy. To this end, we leverage an LLM-guided evolutionary framework that ef- ficiently explores the space of RA functions, yielding Discovered Adversarial Im- itation Learning (DAIL), the first meta-learnt AIL algorithm. Remarkably, DAIL generalises across unseen environments and policy optimization algorithms, out- performing the current state-of-the -hete-rrt of human- designdd baselines. Finally, we analyse why DAIL leads more stable training, offering novel insights into the role of RA functions in the stability of AIL.
GAIL
AIRL
FAIRL
DAIL
10
10
1
1.0
Ass a)
0
0
0.5
1
0.0
-10
10
10
10
10
10
-10
10
log
log
(Arls.
log
log
(Ads..))
(L2S, RT
Figure 1: Visualization of the different reward assignment functions.
INTRODUCTION
Reinforcement Learning (RL) has achieved impressive results across a range of complex do- mains (Mnih et al., 2015; Silver al., 2016), conditioned on the availability of well-defined and in- formative reward functions. However, in many real-world settings, specifying such reward functions
Corresponding author
Research conducted while SMU Research Engineer

---

Published as a conference paper at ICLR 2026
is either prohibitively difficult or entirely infeasible, whereas providing demonstrations of the de- sired behavior is often easier and cost-effective. This motivates the paradigm of Imitation Learning (IL; Argall et al. (2009); Schaal (1999)), which seeks to learn policies directly from expert demon- strations. II. is particularly well-suited for applications such as autonomous driving (Pomerleau, 1988) and robotic manipulation (Argall al,, 2009), where hand-crafting precise reward functions poses a significant challenge.
A particularly effective approach within imitation learning is Adversarial Imitation Learning (AIL; Ho & Ermon (2016)), which draws inspiration from Generative Adversarial Networks (GANs; Goodfellow et al. (2014)) and is recognized for its strong performance when expert demonstra- tions are limited. Similar to GANs, AlL formulates the learning process as a two-player adversarial game between a generator (i.e., the policy network) and a discriminator network. The policy aims to generate trajectories that are indistinguishable from those of the expert, while the discriminator learns to distinguish between expert- and policy-generated trajectories.
Since it shares similar objective to GANs, AIL inherits some of the training challenges associated with adversarial methods (Arjovsky et al., 2017) notably the issues related to instability. In this work, we focus on one critical factor that underpins stable and effective training: the quality of the learning signal. Previous research in GANs (Goodfellow et al., 2014; Arjovsky et al., 2017) has shown that providing strong and informative gradient signals the generator is essential for improving its performance and ensuring convergence of the adversarial game. In the context of AIL., this learning signal manifests the rewards given to states-action pairs visited by the policy. Although recent non-adversarial approaches (e.g., Kostrikov et al. (2020) and Garg et al. (2021)) have shown promise, their empirical performance remains mixed (Jain al., 2024; Lai al., 2024), and they offer limited flexibility (e.g. for reward shaping (Sapora et al., 2024)), motivating the need to stabilize adversarial imitation learning through more informative reward signals for policy optimization.
We first introduce AIL and reward assignment through the lens of divergence minimization between expert and policy occupancy measures as highlighted in Ghasemipour et al. (2020). This perspective reveals a natural two-stage decomposition of the reward assignment process: (a) Density Ratio (DR) estimation, where the ratio of occupancy measures for each state-action pair is estimated, and (b) Reward Assignment (RA), which maps this ratio to scalar rewards for policy optimization. While prior work has improved stage (a) through better stabilization of the discriminator training (Luo et al., 2024; Wang et al., 2024; Lai et al., 2024) (b) has attracted considerably less attention (Fu al., 2018: 8; Ghasemipour et al., 2020). Following Ghasemipour
al. (2020), we highlight how RA functions (Figure 1) influence policy-learning dynamics in adversarial training, and propose an LLM--based meta learning framework to discover reward assignment functions for improved performance.
This optimization produces RA function that, when integrated into the AIL framework, results in Discovered Adversarial Imitation Learning (DAIL) (Figure 1). When evaluated on unseen environ- ments from the Brax (Freeman et al., 2021) and Minatar ( (Young & Tian, 2019) suites, DAIL out- performs state-of- -oe--o-tthherrrr art baselines, including GAIL (Ho & Ermon, 2016 6). AIRL (Fu et al., 2018), FAIRL (Chasemipour et al., 2020) and - (Orsini et al., 2021). To the best of our knowledge, DAIL is the first meta-learned All. algorithm. We further demonstrate that it also gen- cralizes to policy optimization algorithms not seen during discovery. Finally, by examining DAIL's training dynamics, we show that DAIL enhances performance by producing more informative learn- ing signals.
2 RELATED WORK
Learning from Demonstrations In settings where reward design challenging or limited in ex- pressivity (Argall et al., 2009; Teoh et al., 2025), or where exploration is difficult (Nair et al., 2018), learning directly from expert demonstrations offers a compelling alternative. The simplest approach, Behavior Cloning (BC; Pomerleau (1988); Torabi et al. (2018)), treats imitation supervised learn- ing but struggles in low-data regimes due to compounding errors outside the training distribution. In contrast, distribution distribution-matching methods explicitly align the expert and policy state-action distribu- tions, which helps mitigate distributional shift and improves robustness. Among them, Adversarial Imitation Learning (Ho & Ermon, 2016; Fu et al., 2018; Ghasemipour et al., 2020; Orsini et al.,
2

---

Published as a conference paper at ICLR 2026
2021)-inspire inspired by GANs- has shown strong results but suffers from instability. Previous work has focused on stabilizing the discriminator using improved loss functions (Luo et al., 2024; Wang et al., 2024), architectures (Lai et al., 2024), and regularizers (Orsini et al., 2021; Xiao et al., 2019). However, reward assignment-how discriminator logits should be mapped to rewards-has received little attention. Fu et al. (2018) first highlighted the impact of RA, showing that performance can vary with the structure of the underlying MDP, and proposed methods such as incorporating absorbing states mitigate this dependence. Ghasemipour et al. (2020) and Zhang et al. (2020) highlighted the impact of RA functions on policy optimization and along with Ke et al. (2021) unified various RA functions as instances of divergence minimization within the distribution instrbbution-matcingg matching IL framework, while Orsini et al. (2021) empirically benchmark these functions across multiple tasks. This work is, to the best of our knowledge, the first that explores the discovery of novel reward assignment functions, as an alternative to relying on the existing human-derived ones. Recent non-adversarial iistribution-matching IL methods, such as ValueDICE (Kostrikov et al., 2020) and IQ-Leamn (Garg et al., 2021), have demonstrated promise but showed mixed empirical performance, with AIL vari- ants often performing comparably or better across multiple benchmarks (Lai et al., 2024; Jain et al,, 2024). Additionally, unlike AIL approaches, they provide limited flexibility in accommodating scenarios such as state-only demonstrations (Torabi et al., 2018; Jain et al., 2024) and reward shap- ing (Sapora et al., 2024). Another closely related area is Inverse Reinforcement Learning (Ng & Russell, 2000), which seeks to infer the underlying reward function that best explains the expert pol- icy (Skalse & Abate, 2023; Chirra et al., 2024). This contrasts with our narrower focus on directly optimizing imitation policies and designing reward assignment functions that ensure their stable training. On the other hand, works such as (Ma et al., 2024a;b; Li al., 2025) directly tackle the problem of reward design in MDPs by leveraging LLMM-based evolution, ultimately outperforming human-designed rewards.
Meta Learning Meta-Learning (also known as 'learning to learn') aims to automatically discover learning algorithms through end-to-end optimization (Schmidhuber, 1987; Thrun & Pratt, 1998) and has seen extensive application in RL (Beek et al., 2025). Historically, RL training pipelines have been constrained by CPU-bound simulators. The advent of GPU-accelerated environments, how- ever, has enabled speedups of up 1000x, revitalizing interest in Meta-RI.. This has facilitated the discovery of novel loss functions (Lu et al., 2022; Jackson et al., 2024), activation functions (Nadimpalli et al., 2025), and optimizers (Goldie et al., 2024) for Deep RL. Nevertheless, differen- tiating through the learning process of an RL algorithm remains a challenging problem, motivating the use of black-box optimization methods (Salimans et al., 2017; Chen et al., 2023b). However, these approaches typically suffer from high sample complexity and limited interpretability. Large language models (L.LMs) provide a promising alternative as their broad domain knowledge and code-generation capabilities can enable more interpretable and effective learning algorithms (Goldie et al., 2025). Beyond RI., LLMM-based black box optimization has demonstrated success wide range of domains, including environment generation (Faldor et al., 2024), neural architecture search (Chen et al., 2023a), combinatorial optimization (Ye et al., 2024), mathematics (Romera- Pareedes et al., 2024), and more (Novikov et al., 2025).
3 BACKGROUND
Preliminaries We consider a Markov Decision Process (MDP) M defined by the tuple (8, A, P,r,y, n). where S denotes the set of states, A is the set of actions, a(s'Is,a) E [0, 1] is the transition probability, r(s, a) E R is the reward function, tor, and μas) E A(S) the initial state distribution. A policy π (-1s) E A(A) is a distribution over the set of valid actions for state
E [0, 1] is the discount fac-
A trajectory т = {(s,, } denotes the state-action pairs encountered by executing π in M. For a given policy म, the occupancy measure p= a) is defined as (s) = (1 Br (E 7 PP(s,
8,a, a)]. Intuitively, it can be in- terpreted as the distribution over state-action pairs that the agent encounters while following pol- icy æ. A one-to-one correspondence exists between π and Pa (Syed et al., 2008), allowing use them interchangeably. For functions f(s, a) dependent only on state-action pairs, we have E,er Ee= 7 [st, at] - ((sa))p> [f a)]. We leverage this identity to use the two pectations interchangeably when optimizing over f.
to
3

---

Published as a conference paper at ICLR 2026
Distribution Matching IL The goal of iettributton-matccinn imitation learning is to find a policy '² whose occupancy measure Pr closely aligns with that of the expert Pg. In this paper, we focus on f-divergence minimization unifying framework encompassing many distribution-matciing II. algorithms (Ghasemipour et al., 2020). Formally,
n = arg min Dy(PEllea),
where the f-divergence D is defined as
Dr(pellp*) = B(s,))ppp [
Pr [
(2)
with f: R+ - R convex and satisfying f(1) = O. The ratio CE, known as the density ratio, is formally the Radon-Nikodym derivative (Halmos, 1950), which quantifies the pointwise discrepancy between pr and Px. In practice, we only have access finite set of expert demonstrations rather than the true occupancy measure PF.
Adversarial Imitation Learning An adversarial approach solving Equation I involves the fol- lowing iterative steps: (1) Policy Rollouts: Generate trajectories by executing the current policy m. (2) Density Ratio Estimation: Multiple approaches have been proposed to estimate the density ratio P. (Wang et al., 2024; Lai et al., 2024). In this work, we adopt the most common approach (Orsini et al.,, 2021) of training classifier (diseriminator) distinguish expert from policy-generater (s, a) pairs via binary cross-entropy loss. An optimal discriminator'' 's logits (pre-softmax) would then correspond to log Pe (refer to Appendix A. for derivation) and (3) Reward Assignment and Policy Improvement: Depending on the f-divergence being minimized, each (s, a) pair visited by receives reward r(s,a) ry pa (E( ), where Ty R4 + E is defined as the reward assignment function. Table summarizes common divergences and their corresponding reward as- signment functions. The rewards are then used to update the policy, and Steps 1-3 are repeated until convergence.
Table I: Reward assignment functions for different f-divergences, where & é= log pF(s.a) (s.a)
p=(s.))
Divergence
Algorithm
Reward Assignment Function
Forward KI.
FAIRL (Ghasemipour et al., 2020)
(8.a)) { ((s a) softplus (((s a))
Backward KL
AIRL (Fu al., 2018)
Jensen-Shanne
GAIL (Ho & Ermon, 2016)
Unnamed fddiv
GAlL-heuristic (Orsini et al., 2021)
(-((s))) -sofpplus(((ssa))
softpluus
4 PROBLEM DEFINITION
Adversarial methods are often considered unstable due to their reliance optimizing min-max objectives, akin to GANs (Goodfellow et al., 2014). To mitigate this instability, prior work has predominantly focused on Step 2 by improving discriminator training (Luo et al., 2024; Wang al., 2024).
In this paper, focus Step 3, highlighting that providing informative learning signal is crucial for effective policy improvement and overall adversarial training. Originally discussed by Ghasemipour et al. (2020), Figure 1 illustrates how the reward assignment function shapes the pol- icy's learning dynamics. The AIRL RA function encourages the policy to visit state-action pairs where expert visitation exceeds its own and penalizes it equally when it surpasses the expert. In contrast, the GAIL RA function only incentivizes matching the expert on underrepresented pairs, while its heuristic variant does the opposite. FAIRL employs a more nuanced approach: rewards the policy for slightly exceeding expert visitation, but imposes steep penalties when the expert dom- inates. Ghasemipour al. (2020) argues that this drives the policy to gradually expand and cover
While the original convex function f can be applied directly, most approaches utilize the function derived from its variational representation, detailed in the Appendix A
4

---

Published as a conference paper at ICLR 2026
Generated Mutation
Eoergesze (ongta)))
Hiclegtngy reEm a'bppaacttennDDagtt
reant Pecteetes recure tannCoggtaa)
Population
Stvrggeees
Frwird Pow
Inner Training Loop
Expert Demonstrations
Rewerd Funcsion
t
lisay cross
Finasa
Diseriminator
loss
Top Candidates
Enviromennt
comergence
octiin
Figure 2: Visualization of the LLM-guided evolution. Appendix B contains the pseudocode of the framework.
the expert distribution from the outside in. Furthermore, although existing RA functions (Table 1) are derived from well-established f-divergence theory, they neglect the practical stability challenges that arise during training.
This raises the question: Can we meta-learn a reward assignment function that results in stable and effective adversarial training? To evaluate training quality, we measure the divergence between the expert and the policy after training. Specifically, we use the Wasserstein distance (Rubner et al., 1998), computed between rollouts generated by the policy and expert demonstrations. We chose this metric due to its robustness and sensitivity in measuring distances between occupancy distributions in RL settings (Luo al., 2023; Rupf et al., 2024).
4.1 FORMAL DEFINITION
We formalize the meta-learning problem of discovering RA functions
min W(PEs s.t. π T' = arg max (pell(p
(3)
where W denotes the Wasserstein distance, and n" is the optimal policy obtained by iterating over Steps 1-3 with reward assignment function ry- We remove additional constraints on r'y (such as convexity), to enable the exploration of more expressive RA functions beyond those derivable from classical f-divergences. While this approach foregoes theoretical convergence guarantees, the trade- off is justified by the empirical feedback that ry receives with policy training. As our results show, the discovered reward assignment functions exhibit robust generalization properties.
DISCOVERING RA FUNCTIONS VIA EVOLUTIONARY SEARCH
Optimizing Equation 3 is challenging because it requires backpropagating gradients through the entire adversarial training loop, which is generally computationally intractable. Consequently, prior work has typically relied on black-ba box methods for such bilevel optimization problems (Goldie et al., 2024; Lu al., 2022). In this work, adopt LLM-guided evolutionary framework that has been shown to be sample-efficiens interpretable, and capable of discovering generalizable algorithms in meta-RI. (Goldie al., 2025).
5

---

Published as a conference paper at ICLR 2026
The RA function ry is represented directly as code, enabling expressive and interpretable formula- tions. To evolve new candidate functions, we prompt the LLM to intelligently combine and mutate parent programs, guided by their structural and behavioral characteristics. LLMs particularly well-suited for this setting for two key reasons: (1) code (Python) is Turing-complete (Faldor et al., 2024), allowing the search space to encompass rich class of reward assignment functions and (2) the pretrained knowledge encoded in LLMs provides a strong inductive bias, helping to navigate the vast search space more effectively.
Next, we describe the LLM-guided search algorithm used in our work. Importantly, our main con- tribution lies not in the evolutionary algorithm itself, but the formulation and optimization of the meta-lcarning objective (Eq. 3). LLM-guided black-box optimization does not follow rigid standard-many approaches share a common structure but differ in finer details. For completeness, we outline the variant adopted in this work, which closely resembles EvolveAct (Nadimpalli et al., 2025), Evo-proe evo-prompting (Chen al., 2023a), and FunSearch Coomera-Paredee et al., 2024). More so- phisticated strategies (Fernando et al., 2024; Novikov et al., 2025) are complementary to this work.
5.1 COMPONENTS OF THE EVOLUTIONARY SEARCH
Base Population. We initialize the search with reward assignment functions from estab- lished f-divergences--Gar
Normalln ized 79 Distance
Normalized Return
FAIRL, AIRL,
0.98
and GAlL-heuristics (Table 1)-to provide robust foundation for the evolutionary search.
0.88
0.96
0.85
Fitness Evaluation. Each candidate function Ty is evaluated by training a policy to conver- gence using ry the reward assignment func- tion, then measuring the Wasserstein distance between its rollouts and the expert's. This score serves as the fitness criterion for selection.
0.63
0.92
0.80
Generation
Figure 3: Performance across generations on the Minatar Spacelnvaders environment. We report the best-performing member per generation, with Generation 0 denoting the base population. W distance is normalized relative to the best base member (GAIL).
Crossover. To generate new candidate func- tions, we sample parent pairs {r,, r2) from the current population and pass them the LLM along with their fitness scores. The LLM then prompted synthesize new func- tion rys that combines desirable properties of the parents, with the goal of improving perfor- mance (e.g., by blending their functional forms). The detailed prompt format and representative examples of generated candidates are provided in Appendix C.
5.2 SEARCH PROCEDURE
The LLM-guided evolutionary framework unfolds over multiple iterations as follows:
Initial Population: Initialize the search with a base population of reward assignment functions derived from known f-divergence formulations.
Iterative Evolution:
At each generation: (1) Randomly sample M pairs of reward assignment functions from the current population. (2) For each pair, use the LLM to generate N new candidate functions by recombining and refining the parent functions.
(3) Evaluate all M N candidates using the ciarrbution-matchhig fitness score. (4) Select the top K candidates to populate the next generation.
Termination: Repeat until stopping criterion is achieved (c.g., performance plateau).
6 EMPIRICAL STUDIES
We begin by describing the experimental setup, after which we present the results. We conducted our experiments two benchmark suites: MuloCo control tasks (Ani, Reacher, Walker2d, HalfChee-

---

Published as a conference paper at ICLR 2026
Brax
Median
Optimality Gap
GAIL-h FAIRL (FKL)
AIRL (RKL)
GAIL (J5)
0.25
0.50
0.75
0.2
0.4 0.6
0.0
0.3
0.9
0.25
0.50
0.75 1.00
MinAtar
Median
IQM
Mean
Optimality Gap
GAIL-h
FAIRL (FKL)
AIRL (RKL)
GAIL (JS)
DAIL
0.00
0.25 0.50 0.75
0.00
0.25
0.50
0.75
0.00
0.25
0.50
0.75
0.25
0.50
0.75
1.00
Normalized Returns
Figure 4: Aggregate performance on the Brax and Minatar suites (excluding SpaceInvaders),
tah, and Hopper) (Todorov et al., 2012) and Minatar (Asterix, Spacelnvaders, and Breakout) (Young & Tian, 2019). For each task, collect 10 successful expert demonstrations from PPO-trained policy (Schulman et al., 2017) and subsample every 20th transition, following standard practice (Ho & Ermon, 2016). Unless stated otherwise, we optimize the policy using PPO and regularize the discriminator with a gradient penalty (Gulrajani et al., 2017). Consistent with prior work in II, adopt fixed-length episodes of 1000 timesteps for MuJoCo environments (Gleave et al., 2022). However, we do not enforce this for the Minatar tasks since doing so significantly degraded perfor- mance. Our implementation fully written in JAX (Bradbury et al., 2018), using PureJaxRL (Lu et al., 2022), Brax (Freeman et al., 2021) for MuJoCo environments, Gymnax (Lange, 2022) for Mi- natar, and OTT-JXX (Cuturi et al., 2022) for computing Wasserstein distance. The complete training hyperparameters are provided in Appendix D. We normalize returns using min-max scaling between random and expert policy performance, with all results averaged over 16 independent seeds.
6.1 EVOLUTIONARY SEARCH
We performed the evolutionary search on the Minatar Spacelnvaders environment, which has pre- viously been shown to facilitate the discovery of generalizable meta-RL algorithms (Jackson al., 2024). For the search, we GPT-4.1-mini, selected for its strong performance-cos tradeoff. Throughout evolution, we the PPO and discriminator hyperparameters, evaluating 200 candi- date RA functions over span of three hours. Full details of the evolutionary hyperparameters provided in Appendix D.
The evolutionary trajectory is shown in Figure 3. The best-performing reward assignment function discovered at the end of the search is:
ddise (x) = 0.5 sigmoid(2) [tanh(x)) + 1].
(4)
Building on this, we introduce Discovered Adversarial Imitation Learning (DAIL), which applies standard imitation learning loop with Tdise as the reward assignment function. Remarkably, DAIL reduces the W-distance to expert trajectories by 20 % and improves normalized returns by 12.5% compared to the best baseline (GAIL).
6.2 GENERALIZATION
We now turn to the out-of-distribution performance of DAIL. As first step, we benchmark our method against previous RA functions (Table 1) in Brax and Minatar environments (excluding Mi- natar SpaceInvaders). All methods share identical hyperparameters, differing only in the choice of RA function. Aggregated scores using the x rliable library (Agarwal et al., 2021) are shown in Figure 4. Among the baselines, AIRL performs slightly better than GAIL Brax but performs

---

Published as a conference paper at ICLR 2026
Algorithm X
Algorithm Y
0.7
DAIL
DAIL
FAIRL (FKL)
Normalized Return
0.6
GAIL
0.5
DAIL
GAIL-h
0.4
DAIL
GAIL (S)
0.3
0.2
DAIL
AIRL (RKL)
0.1
0.0
PIX
0M
2M
8M
10M
Steps
Figure 5: (Left) Probability of improvement of DAIL over baselines on Brax. (Right) Performance comparison between DAIL and GAIL (on Minatar SpaceInvaders) using A2C. We report the mean and standard error (SEM).
poorly on Minatar, similar to GAIL-heuristic, This is likely because their reward assignment func- tions yield predominantly negative rewards, incentivizing agents terminate episodes early. FAIRL performs poorly across both suites, likely due to its exponential, unbounded reward decay for posi- tive log-ratios, which destabilizes training. Overall, the baseline trends align with the observations from (Orsini al., 2021).
Figure demonstrates DAIL's effectiveness across both benchmark suites. On Minatar, DAIL sig- nificantly outperforms all baselines across all evaluation metrics. On Brax, DAIL outperforms base- lines on most metrics, with slightly lower median than AIRL and statistically significant gains in mean performance. To quantify the robustness of these improvements probability of improvement (Agarwal et al., 2021) metric, which estimates the probability that DAIL outperforms baseline a probability improvement greater than 0.5 against all baselines, providing further evidence of its superior performance. To assess DAIL's generalization beyond the policy optimizer (PPO) used during evolution, we evaluate its performance with A2C (Mnih et al., 2016). As shown in Fig- ure 5 (right), DAIL maintains significant performance advantages over GAIL showcasing its effec- tiveness across different policy optimization algorithms. Finally, we assess DAIL's generalization performance under various discriminator regularization strategies proposed by Orsini et (2021), including weight decay, additional entropy bonus, and spectral normalization of the discrimi- nator weights. We also consider the case without any regularization. shown in Table 2, DAIL. outperforms GAIL in out of 5 regularization regimes.
Brax, we employ the
randomly chosen task. As shown in Figure 5 (left), DAIL achieves
Algo
Env
none
w-decay
entropy
spectral
grad-pen
Asterix
0.88 = 0.03
1.33 0.3
0.12 + am
0.92 +: 0.
0.66 0.0
Breakout
0.81 0.07
0.74
0.91 0.0
0.77 007
1.01 8.0
DAIL
SpaceInvaders
0.71
0.81
0.80
0.70
0.90
Overall
0.80 +: 0.00
0.96 0.03
0.61 0.01
0.80 + 0.04
0.85 0.01
Asterix
1.18 + 0.03
1.44
0.48
0.22 + 0.03
0.52
Breakout
0.76 + 0.00
0.52 + 0.10
0.89 0.01
0.33 + 0.10
0.85 000
GAIL
SpaceInvaders
0.61 + 0.09
0.34 + 0.09
0.81 + 0.00
0.42 4 0.08
0.81 + 0.00
Overall
0.85 0.04
0.76
0.73 0.01
0.32 + 0.65
0.73
Table 2: Performance of DAIL and GAIL under different discriminator regularization strategies. Hyperparameters are adopted from Orsini et al. (2021). Reported values denote the mean and stan- dard error across runs.
8

---

Published as a conference paper at ICLR 2026
1.5
8
PPO
GAIL
GAIL
1.2
Policy Entropy
DAIL
DAIL
4
Reward
1.0
2
0.8
0.5
-2
0.2
500
1000
1500
2000
2500
0.0
-3
2
-1 Log-Ratio
Updates
Figure 6: (Left) Policy entropy during training (on HallCheetah) with different RA functions; PPO with simulator rewards is presented as reference. Results show mean SEM. (Right) Distribution of log-density ratios during training on Minatar SpaceInvaders, estimated via kernel density fitting.
6.3 ANALYSIS
Next, we investigate why DAIL leads to such strong performance. As shown in Figure 1, Taisc exhibits an S-shape with a sharper gradient and a slight rightward shift compared to a standard sigmoid. Importantly, it is bounded within the interval [0, 1]. unlike existing baselines. Prior work has shown that bounding rewards can stabilize Deep RL (Mnih et al., 2015; Van Hasselt et al., 2016), which we hypothesize contributes to DAIL's performance.
We assess training stability by tracking policy entropy and comparing it to a PPO agent with ac- cess to simulator rewards. As shown in Figure 6 (left), policies trained with raise converge to lower entropy, closely matching the PPO baseline. This suggests that Vdise delivers a rich and informa- tive signal, enabling sharper action distributions that reflect confident behavior and effective reward maximization. In contrast, GAIL's RA function produces noisier signals, leading to higher-entropy policies and greater uncertainty in action selection.
To gain deeper insight, we examine the distribution of log density ratios log 2s of state-action pairs visited during training. We compare the distributions between DAIL and GAIL and analyze the interaction with their respective RA functions. Figure (right) shows that a large fraction of the log- ratios lie within the interval [- 1, 0] for both methods, with a long tail extending to approximately -2 -2-a region we identify as indicative of random policy behavior. Tdise saturates near zero for x S -188, effectively filtering out noisy low-quality state-action pairs while maintaining informative gradients for moderately performing ones. In contrast, GAIL's reward function assigns high positive values even at
p.
2, thereby rewarding stateacctions pairs corresponding near-randee candrrnnom policies. We posit that this over-sensitivity to low-quality behavior contributes to the noisier reward signals and instability observed in GAIL's training dynamics. Note that the findings in Figure generalize across all test environments; results for which are omitted due to space constraints.
To further test this hypothesis, conduct ablation study the individual components of rdiscs comparing it against sigmoid(z) and 0.5 [tanh(*)) + 1]. All three functions map to [0.1] and exhibit S-shaped curves, but differ in their response characteristics: raise closely follows the tail of the density ratio distribution, while the other two functions provide noisier reward signals that remain positive around r -1.8, with sigmoid((z) being the least responsive due to its relatively flat profile. These differences are validated empirically (Figure 7), where Taise achieves the best performance, followed by 0.5 [tanh(x) 1), with sigmoid( performing worst.
6.4 STABILITY
To assess the stability of the evolutionary process, conduct additional independent evolution- ary run on the Minatar Spacelnvaders environment. The top-5 RA functions discovered in both runs are presented in Table 8 and the top-3 are plotted in Figure 8. We that the 6 evolved RA functions exhibit highly similar structures. Further, they maintain informative gradients within the range [- and saturate rapidly thereafter, consistent with analysis in Section 6.3. Further-

---

Published as a conference paper at ICLR 2026
1.5
1.2
sigmoid
Taus
Normalized Return
0.8
Reward
1.0
0.5 [tanh +
0.6
0.8
Fdtec
0.2
0.5 [tanh + 1]
sigmoid
0.55
0.0
-2.0
-1.0
0.5
10M
Log-Ratio
Steps
Figure 7: (Left) Log-density ratio distributions during training (Minatar SpaceInvaders), estimated via kernel density estimation. (Right) Idkac vs. component function performance (Minatar SpaceIn- vaders).
more, we conduct evolutionary run in the MinAtar Breakout environment and find that DAIL emerges among the top-5 RA functions in the final generation, demonstrating both stability of the evolutionary process and the effectiveness of DAIL.
7 CONCLUSION
Summary This work highlights the impor- tance of the RA functions influencing both policy optimization and the overall stability of All-ann aspect that has received relatively lit- tle attention. We introduce novel approach using LLM-guided evolutionary search to auto- matically discover optimal reward assignment functions, resulting in DAIL, the first meta- learned AlL algorithm. Experimental valida- tion demonstrates DAIL's superior and consis- tent performance across unseen environments and policy optimization algorithms. Through analysis of DAIL's discovered reward function Tdise and its impact training dynamics, we provide novel insights into these performance improvements.
Run 1-
Run 1-
Run
Reward
Run #1
Run 2 #2
0.5
Run #3
GAIL
0.0
-3
-1 Log-Ratio
Figure Comparing the top-3 discovered RA functions across two independant evolutionary runs. Note that Run #1 corresponds DAIL.
Limitations and Future Work While DAIL
demonstrates strong generalization, the discovered RA function raise does not correspond valid f-divergence and therefore lacks theoretical guarantees. Moreover, despite its strong performance, the RA function of DAIL remains static throughout training and does not adapt the training state (e.g., number of updates remaining, loss, observed log-ratios). Exploring time-aware RA functions those that condition on the training state state- could yield richer, more informative learning signals, similar to Jackson et al. (2024). Additionally, including more information into the LLM's context such as environment information may facilitate more effective crossovers.
Finally, it would be valuable to evaluate DAIL's generalization to more complex benchmarks, such as the Atari-57 (Bellemare et al., 2013) and Procgen (Cobbe et al., 2020) suites. Prior work in meta- learning Oh et al. (2020) has demonstrated that training across diverse set of environments yields more robust algorithms. For instance, Oh et al. (2025) introduced the Disco57 and Discol03 suites, consisting of 57 and 103 environments, respectively. Investigating the discovery of AIL algorithms meta-trand trainer on such large-scale environment collections is an exciting direction for future research, though it currently lies beyond our computational budget.
10

---

Published as a conference paper at ICLR 2026
8 ETHICS STATEMENT
As with other meta-learning approaches, the automated discovery of the algorithms (DAIL) obscures its properties and behavior, making analyses like those in Section 6.3 both challenging and essential for understanding such algorithms. Additionally, as with other IL algorithms, these advances hold promise for safer and more capable Al systems but also introduce risks, including misuse (e.g.. imitation of harmful behaviors) and bias (inherited from expert data).
REPRODUCIBILITY STATEMENT
All hyperparameters and experimental details are reported in Section 6 and Appendix D, with com- plete information prompts and the LLM in Appendix C. The supplementary material provides code (including plotting scripts) and evaluation data to fully reproduce all the main results in this work.
10 ACKNOWLEDGMENTS
This research/project was supported by the Singapore Ministry of Education (MOE) Aca- demic Research Fund (AcRF) Tier grant (Proposal ID: 24-S1S-SMU-010).
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances in neural informa- tion processing systems, 34:29304-29320, 2021.
Brenna D. Argall, Sonia Chernova, Manuela Veloso, and Brett Browning, A survey of robot learning from demonstration. Robot. Auton. Syst., 97(5):469-488) May 2009. ISSN 0921-8890. doi: 10. 10.100 wtwp::o..00.....00000 10.024.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 214-223.
PMLR, 06-11 1 Aug 2017.
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shi- mon Whiteson. A tutorial on meta-reinforcemmet learning. Foundations and Trends@ in Machine Learning, 18(2-3):224-384 2025. ISSN 1935-8245. doi: 10.1561/2200000080.
Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning envi- ronment: evaluation platform for general agents. J. Artif. Int. Res., 47(1):253-27, May 2013.
ISSN 1076-9757.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPPy programs, 2018.
Angelica Chen, David Dohan, and David So. Evoprompting: Language models for code-level neural architecture search. Advances in neural information processing systems, 36:7787-7817, 2023a.
Xiangning Chen, Chen Liang. Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, and Quoc V Le. Symbolic discovery of optimization algorithms. In Thirty-seventh Conference on Neural Information Processing Systems, 2023b.
Shashank Reddy Chirra, Pradeep Varakantham, and Praveen Paruchuri. Preserving the privacy of reward functions in mdps through deception. arXiv preprint arXiv:2407.09889, 2024.
Karl Cobbe. Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural genera- tion to benchmark reinforcement learning. Proceedings of the 37th International Conference Machine Learning, ICML'20. JML.. 2020.

---

Published as a conference paper at ICLR 2026
Marco Cuturi, Laetitia Panngaapaxan,hos Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint 1rXIv:2201.12324, 2022.
Maxence Faldor, Jenny Zhang, Antoine Cully, and Jeff Clune. Omni-epic: Open-endedness via mod- cls of human notions of interestingness with environments programmed in code. arXiv preprint arXiv:2405. 15568, 2024.
Chrisantha Fernando, Dylan Banarse, Henryk Michalewski, Simon Osindero, and Tim Rocktäschel. Promptbreeder: self-referential self-improvement via prompt evolution. In Proceedings of the 41st International Conference on Machine Learning, ICML'24. JMLR.org, 2024.
C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax a differentiable physics engine for large scale rigid body simulation, 2021.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adverserial inverse rein- forcement learning. In International Conference Learning Representations, 2018.
Divyansh Garg, Shuvam Chakraborty, Chris Cundy, Jiaming Song, and Stefano Ermon. Iq-learn: Inverse soft-q learning for imitation. Advances in Neural Information Processing Systems, 34: 4028-4039, 2021.
Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu. A divergence minimization perspective on imitation learning methods. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura (eds.), Proceedings of the Conference on Robot Learning, volume 100 of Proceedings of PMLR, 30 Oct-0l Nov 2020.
Machine Learning Research, pp.
Adam Gleave, Mohammad Taufeeque, Juan Rocamonde, Erik Jenner, Steven Wang. Sam Toyer, Maximilian Ernestus, Nora Belrose, Scott Emmons, and Stuart Russell. imitation: Clean imitation learning implementations. arXiv preprint arXiv:2211.1197, 2022.
Alexander D Goldie, Chris Lu, Matthew T Jackson, Shimon Whiteson, and Jakob Foerster. Can learned optimization make reinforcement learning less difficult? Advances in Neural Information Processing Systems, 37:5454-5497, 2024.
Alexander David Goldie, Zilin Wang, Jakob Nicolaus Foerster, and Shimon Whiteson. How should we meta-learn reinforcement learning algorithms? In Reinforcement Learning Conference, 2025.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger (eds.), Advances Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Ishaan Guirajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Im- proved training of wasserstein gans. In Proceedings of the 31st International Conference Neural Information Processing Systems, NIPS' 17, pp. 5769-5779, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Paul R. Halmos. Measure Theory. Springer, New York, 1950. ISBN 9780387900888.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In D. Lee, M. Sugiyama, U. Luxburg. I. Guyon, and R. Gareett (eds.), Advances volume 29, Curran Associates, Inc., 2016.
Advances in Neural Information Processing Systems,
Matthew Thomas Jackson, Chris Lu, Louis Kirsch, Robert Tjarko Lange, Shimon Whiteson, and Jakob Nicolaus Foerster. Discovering temporally reinforcement leamiing algorithms. In The Twelfth International Conference on Leurning Representations, 2024.
Arnav Kumar Jain, Harley Wiltzer, Jesse Farebrother, Irina Rish, Glen Berseth, and Sanjiban Choud- hury. Non-adversarial inverse reinforcement learning via successor feature matching. arXiv preprint arXiv:2411.00 1.07007, 2024.
12

---

Published as a conference paper at ICLR 2026
Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srinivasa. Imitation learning as f-divergence minimization. In Steven M. LaValle, Ming Lin, Timo Ojala, Dylan Shell, and Jingjin Yu (eds.), Algorihhiic Foundations of Robotics X/V, pp. 313-329, Cham, 2021. Springer International Publishing. ISBN 978-3-030-66673--8.
Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Imitation learning via off-policy distribution matching. In International Conference on Learning Representations, 2020.
Chun-Mao Lai, Hsiang- Chan Wang, Ping-Chure Hsich, Yu-Chiang Frank Wang, Min-Hung Chen, and Shao-Hua Sun. Diffusion-reward adversarial imitation learning. The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
Robert Tjarko Lange. gymnax: A JAX-based reinforcement learning environment library, 2022.
Pengyi Li. Jianye HAO, Hongyao Tang, Yifu Yuan. Jinbin Qiao, Zibin Dong. and YAN ZHENG. R*: Efficient reward design via reward structure evolution and parameter aligmment optimization with large language models. In Forty-second International Conference on Machine Learning, 2025.
Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems, 35:1 35:16455- 16468, 2022.
Tianjiao Luo. Tim Pearce, Huayu Chen, Jianfei Chen, and Jun Zhu. C-GAIL: Stabilizing generative adversarial imitation learning with control theory. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Yicheng Luo. Zhengyao Jiang, Samuel Cohen, Edward Grefenstette, and Marc Peter Deisenroth. Optimal transport for offline imitation learning. arXiv preprint arXiv:2303 13397111,,1 2023.
Yecheng Jason Ma, William Liang, Guanzhi Wang, De-An Huang, Osbert Bastani, Dinesh Jayara- man, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Eureka: Human-level reward design via coding large language models. In The Twelfth International Conference on Learning Representa- tions, 2024a.
Yecheng Jason Ma, William Liang, Hung-Ju Wang, Sam Wang, Yuke Zhu, Linxi Fan, Osbert Bas- tani, and Dinesh Jayaraman. Dreureka: Language model guided sim-to-real transfer. arXiv preprint wwXxv:2406.01966 2024b.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonogiou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540);529 533, 2015. doi: 10.1088/n 14236.
Volodymyr Mnih, Adrià Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning, 2016.
Kalyan Varma Nadimpalli, Shashank Reddy Chirra, Pradeep Varakantham, and Stefan Bauer. Evoly- ing RI: Discovering new activation functions using L.LMs. In Towards Agentic Al for Science: Hypothesis Generation, Comprehension, Quantification, and Validation, 2025.
Ashvin Nair, Bob McGrew, Marcin Andrychowicz. Wojeiech Zaremba, and Pieter Abbeel. Over- coming exploration in reinforcement learning with demonstrations. In 2018 IEEE international conference on robotics and automation (ICRA), pp. 6292-6299. IEEE, 2018.
Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In Proceedings of the Seventeenth International Conference on Machine Learning. ICML '00, pp. 663-670, San Francisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc. ISBN 1558607072.
13

---

Published as a conference paper at ICLR 2026
Alexander Novikov, Ngån Vü, Marvin Eisenberger, Emilien Dupont, Po-Sen Huang, Adam Zsolt Wagner, Sergey Shirobokov, Borislav Kozlovskii, Francisco J. R. Ruiz, Abbas Mehrabian, M. Pawan Kumar, Abigail See, Swarat Chaudhuri, George Holland, Alex Davies, Sebastian Nowozin, Pushmeet Kohli, and Matej Balog. Alphaevolve: coding agent for scientific and algorithmic discovery, 2025.
Junhyuk Oh, Matteo Hessel, Wojeiech M. Czamecki, Zhongwen Xu, Hado van Hasselt, Satinder Singh, and David Silver. Discovering reinforcement learning algorithms. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook,
NY, USA, 2020. Curran Associates Inc. ISBN 518-088-09761111
Junhyuk Oh, Gregory Farquhar, Ila Kemaev, et al. Discovering saate-o--the-aa reinforcement learn- ing algorithms. Nature, 2025. doi:
Manu Orsini, Anton Raichuk, Léonard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthicu Geist, Olivier Bachem, Olivier Pietquin, and Marcin Andrychowicz. What matters for adversarial imitation learning? In Proceedings of the 35th International Conference on Neural Information Processing Systems, NIPS '21, Red Hook, NY, USA, 2021. Curran Associates Inc. ISBN 9781713845393.
Dean A. Pomerleau. Alvinn: An autonomous land vehicle in neural network. In D. Touretzky (ed.), Advances in Neural Information Processing Systems, volume 1. Morgan-Kaufmannt Kaufmann, 1988.
B. Panerr-Paredes, M. Barekatain, A. Novikov, J. Clift, R. McAllister, A. R. Massner, A. W. Senior, Veličković, M. Botvinick, D. Hassabis, and A. A. Alemi. Mathematical discov- eries from program search with large language models. Nature, 625:468-475, 2024. doi: 1011:3//411866-022--922---
Y. Rubner, C. Tomasi, and L.J. Guibas. A metric for distributions with applications to image databases. In Sixth International Conference on Computer Vision (IEEE Cat. No.98C1136271, pp. 59-66996 66, 1998. doi: 10.1 00.119//CCV11998.177700..
Thomas Rupf, Marco Bagatella, Nico GÄärter,, Jonas Frey, and Georg Martius. Zero-shot offline imitation learning via optimal transport. arXiv preprint arXiv:2410.08751, 2024.
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies as scalable alternative to reinforcement learning. arXiv preprint arXiv: 2017.
Silvia Sapora, Gokul Swamy, Chris Lu, Yee Whye Teh, and Jakob Nicolaus Foerster. Evil:
lution strategies for generalisable imitation learning. In Proceedings of the 41st International Conference on Machine Learning, ICML'24. JMLR.org, 2024.
Stefan Schaal. Is imitation learning the route to humanoid robots? Trends in Cognitive Sciences, 3 (6):233-242, 1999. ISSN 1364-6613. doi: https://doi.org/10.. 10160000011//11660110 (399946-666619910337733333
Jurgen Schmidhuber. Evolutionary principles in self-referential learning. on learning now to learn: The meta-mett-mmtta..hhor Diploma thesis, Technische Universitat Munchen, Germany, 14 May 1987.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXivv 2017.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484 489, 2016. doi: m0.t//rr 16961.
Joar Skalse and Alessandro Abate. Misspecification in inverse reinforcement learning. In Proceed- ings of the Thirty-Sevente AAAI Conference on Artificial Intelligence and Thirty-Fifth Conference on Innovative Applications of Artificial Intelligence and Thirteenth Symposium on Educational Advances in Artificial Intelligence, AAAI22/AAAT233AAA AAAI Press, 2023. ISBN 978-
--57735-880-0. doi: 10.
14

---

Published as a conference paper at ICLR 2026
Umar Syed, Michael Bowling, and Robert E. Schapire. Apprenticeship learning using linear pro- gramming. In Proceedings of the 25th International Conference on Machine Learning, ICML '08, pp. 1032-1039, New York, NY, USA. 2008. Association for Computing Machinery. ISBN 9781605582054. doi: 10.1 155/119915.... 39.156.3398866
Jayden Teoh, Pradeep Varakantham, and Peter Vamplew. On generalization across environments in multi-objective reinforcement learning, 2025. URL /tRL:/tsss///axx..rr//////////5 00799.
Sebastian Thrun and Lorien Pratt. Learning to learn: introduction and overview, pp. 3-17, Kluwer Academic Publishers, USA, 1998. ISBN 0792380479.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 TEEE/RSJ International Conference Intelligent Robots and Systems, pp. 5026-5033, 2012. doi:
Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. In Proceedings of the 27th International Joint Conference on Artificial Intelligence, IICAI'18, pp. 4950-4957.
AAAI Press, 2018. ISBN 9780999241127.
Hado P Van Hasselt, Arthur Guez, Matteo Hessel, Volodymyr Mnih, and David Silver. Learning values across many orders of magnitude. Advances in neural information processing systems, 29, 2016.
Bingzheng Wang, Guoqiang Wu, Teng Pang, Yan Zhang, and Yilong Yin. Diffail: Diffusion adver- sarial imitation learning. Proceedings of the AAAI Conference on Artificial Intelligence, 38(14); 15447-15455, Mar. 2024. doi:
Huang Xiao, Michael Herman, Joerg Wagner, Sebastian Ziesche, Jalal Etesami, and Thai Hong Linh. Wasserstein adversarial imitation learning. arXiv preprint arXiv:/9066 13, 2019.
Haoran Ye, Jiarui Wang, Zhiguang Cao. Federico Berto, Chuanbo Hua, Haeyeon Kim, Jinkyoo Park, and Guojie Song. Reevo: Large language models as hyper-heuristics with reflective evolution. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024.
Kenny Young and Tian Tian. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv preprint 17X1v:1903.03176, 2019.
Xin Zhang, Yanhua Li, Ziming Zhang, and Zhi-Li Zhang. f-gail: Learning f-divergence for gener- ative adversarial imitation learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 12805-12815. Curran Associates, Inc., 2020.
15

---

Published as a conference paper at ICLR 2026
A ON EDIIVERGENCE MINIMIZATION
We present key preliminary results that will support the derivations in later sections.
A.1 BACKGROUND
Note that the results presented here assume a discounted infinite-horizon setting in a discrete MDP for simplicity, but they can be extended to other settings such as finite-horizon problems or continu- state-action spaces.
Definition (Occupancy Measure) For a policy т,, the occupancy measure Pa is defined as:
p" (s.a) (1 T E = = a)
Lemma 1 (Interchange of Expectations). For any scalar function f S x A - R and discount factor y E [0,
E
Σ (ss. =
1 1-7 (( [f(s,a)]-
(2)
t=0
Proof. Starting from the definition the left-hand side, we have
E
y f((s
ET
Σ
Σ P(st =.aa =.0.=a)f(s.a))
(3)
1==
20
7P((s 8, 8,00 -a)= [(0,0))
Byax
(4)
1=0
Pr(88
f(s,a)
(using Eq. 1)
(5)
1
8,0
B(saa><,, [f(s,a]]
(6)
Lemma 2 (Optimal Discriminator). Let P and Q be two distributions over random variable X.
Consider a discriminator parameterized as D(X) = a (E(X)), where a denotes the sigmoid func- tion and D(X) represents the predicted probability that X is drawn from P. If the discrininator is trained via likelihood maximization (i.e., using binary cross-entropy loss), then the optimal discrim- inator satisfies:
P(X) Q(X)
8" (X) = log
Proof. The discriminator is trained using the binary cross-entropy loss:
LBCE (log DXX) + Ex-q/log(1 D(X)))
(8)
To find the optimal discriminator, we maximize LECF. with respect to D(X) pointwise. Taking the derivative of the pointwise objective and setting it to zero yields:
P(X)
D*(X) P(X) + Q(X)
(9)
Substituting (*(X) (X)() we get:
P(X)
(X))
(10)
-
P(X) + Q(X)
Solving for (X) gives:
P(X)
(x) = log
(1I)
((X))
16

---

Published as a conference paper at ICLR 2026
A.2 FAIRL REWARD ASSIGNMENT FUNCTION
The f-divergence between two distributions P and Q is defined as:
D2(P)QQ Ex-Q (Q(x))
(12)
where f: + R is convex function satisfying f(1) O. The ratio P(X) is referred to the Q()
density ratio, and is formally defined as the Radon-Nikodym derivative (Halmos, 1950) of with respect to Q.
In the context of f-divergence-based imitation learning, the objective is find a policy æ that minimizes the divergence between the expert and policy occupancy measures:
x² = arg min (p@ellem*
(13)
(
arg min B(s,a)=pt
(r(s,a)
(14)
(p0(s.0))
arg min (1 7) Eiaa f (=,
(using Lemma 1)
(15)
= (1 - 7) arg max E (s)
(16)
where the reward function ry is defined as:
(PE(S, ( Px (s,
a)(s,a)= -f =
(17)
This formulation establishes that minimizing an f-divergence is equivalent maximizing the ex- pected return under a reward assignment function ry which is a mapping from the density ratio Ps to scalar reward.
Pa
In the case reverse KL divergence, the corresponding f-function is f(z) log x. Defining the p the resulting reward assignment becomes:
log-density ratio ((s,a) log p=(3,a))
TT-RKL(S, a) - ((.a)) E(s. a)
(18)
which corresponds to the reward used in FAIRL (Ghasemipour al., 2020).
A.3 REWARD ASSIGNMENT FUNCTIONS IN GAIL AND AIRL
The f-divergence admits a variational representation:
(PI()) = sup Ex~ p((()] Ex Ex-xlf(((X))) R
(19)
where f" denotes the convex conjugate of f. defined by
(*( = sup
{wo- - f(v)).
vEdom([)
To understand the structure of the optimal function we consider the pointwise optimization of the P(X) the first-order optimality condition Q(x)
integrand in Eq. equation 19. Letting =XX) and becomes:
V (u)
(20)
Under the assumption that differentiable and strictly convex, the gradient of the convex conju- gate satisfies the inverse relationship:
vf*(u)= (Vf)~ (w).
(21)
17

---

Published as a conference paper at ICLR 2026
Substituting into the optimality condition, obtain:
(f)) (u) =
(22)
7 u = (c)--
(23)
Thus, have,
9" 9'(X) f P(X) Q(x)
(24)
In the context of f-divergence based imitation learning, we have, म² arg min Dy(e
(25)
arg min sup B(s,)) ((s,))ppn. If (g(s,a))
(26)
arg min (,,))=p S' PT Pr(s.aa) (PPrll.aa)
(using Eq. 24)
(27)
pE(s.a)
7) arg max E Fr s where the reward function --vaar defined
(28)
PE pE(s, Par (s,
Tf-var a) = =
(29)
This establishes that minimizing the f-divergence between Ps and pe via its variational formulation equivalent to maximizing the expected return under a reward assignment function defined by ff-ar.
In the case of reverse-KL. divergence, the corresponding f f-fucction is f(x) = x llo x, and thus (x)) = 1 log Plugging this into the reward assignment gives:
( (E((,,)
Tf (s,a)
| log
(r(()))) a)
(30)
PE(s, a)
log
-1
(31)
Px (s,
Ignoring the additive constant, the reward assignment under the reverse-KL. divergence simplifies to:
TJ-va-RKL(s,)) E(s, a)
(32)
which corresponds to the reward used in AIRL (Fu et al., 2018).
In the case of Jensen-Shannon divergence, the corresponding f-function is f(z) -(x + log (##2)= + logæ, z, and its derivative f'(x) log x+1
2x
Hence the variational reward
signment becomes:
(s -log (p (p
(33)
1 Pn (8, ())EE
log 1+ pE(s.
1
(34)
2
(s
Ignoring the additive constant, the reward assignment under the JS divergence simplifies
Js a) = log(1 + (s.a))
(35)
which corresponds to the reward used in GAIL (Ho & Ermon, 2016).
By plugging the f-functions of other commonly used f-divergences into Eq. 29, we can derive their corresponding reward assignment functions, summarized in Table 3. While formal adversarial imitation learning (AIL) methods explicitly employ these divergences, they have been explored in the context of non-adversarial imitation learning algorithms such as IQ-Leam (Garg al., 2021).
18

---

Published as a conference paper at ICLR 2026
Table 3: Reward assignment functions derived from different f f-divergences
Divergence
f(x)
ffvar
Forward KL
log æ
Px.
Reverse KLL
x: log 2:
log 2E
Jensen-Shannon
x log x - (x + 1) log
222
log 1 +
A.
Squared Hellinger
(vz - 1)2
Ps
Pri
Pearson x2
(x 1)²
Pz.
Total Variation
} 1 1|
sign
P
PE
A.4 OPTIMIZATION
In the adversarial imitation learning (AIL) framework, optimizing the f-divergence objective in Eq. 13 corresponds to the following iterative procedure that alternates between training a discrimi- nator and updating the policy.
Assume discriminator the form D(s, a) = (l(ss, a)). where f(s, a) is a learned logit function and (-) denotes the sigmoid function. The goal of the diseriminator is to distinguish between state- action pairs sampled from the expert occupancy measure p and those induced by the current policy æ, denoted p.
Step 1: Discriminator Update. The discriminator is trained by maximizing the binary cross- entropy objective:
D' (*s8a) = arg max E(c.a)~pE [log D(s, + ((e,))~ [log(1 - D(s, a))]
(36)
D
Step 2: Policy Update. The policy is then updated to maximize the expected return, where the reward is derived from the discriminator output via function ff/f a), which typically corre- sponds to variational lower bound on the chosen f-divergence:
arg max (( [15
(37)
where the choice of the reward assignment function f/ depends on the f-divergence used. Iterate between Steps and until convergence.
Convergence. As highlighted in (Ghasemipour et al., 2020), under the assumption that the dis- criminator is optimized to its optimum D' at each iteration, the overall procedure converges to a fixed point where the occupancy measure of the learned policy matches that of the expert, i.e., Pa PE.
B PSEUDO-CODE
The complete pseudo-code for f-AIL is provided in Algorithm 1, while the LLM LLM-guided evolution- ary search detailed Algorithm 2.
19

---

Published as a conference paper at ICLR 2026
Algorithm f-AIL: Adversarial Imitation Learning via f-Divergence Minimization
Require: Expert trajectories TE ~ Tg; initial policy, discriminator parameters lo. do: A entropy coefficient, number of updates T
1: for iteration i 0, .,..., T do
2: Sample trajectories T; THe
3: Update discriminator parameters Φ; 0541 using the gradient:
& (s,,))=T [log (s, a)] + ((s,)) [log(1 Do a))
4: Construct rewards using the discriminator logit: r(s. a) + (y((,,_,1 (s.
5:
Compute advantage estimates A(s, a) from trajectories т; using r(s,
6:
Update policy parameters 0: -> 0.+1 via PPO by optimizing
Vollr [min (r. (r) a), clip((:(())
A) a))] XV (v
where (0)
To s) ne, [a
is the likelihood ratio, and H. is the causal entropy.
7:
end for
8:
Sample trajectories Ter
9:
Calculate Wasserstein distance between occupancy measures:
D wassecsstin -W((s, a) ((,(((s))) ()x)) where (s, a)xx ~ 7 and (s, a) ~ TE
10: return
Tg. D
wasserstein
Algorithm LLLM-Guidet Evolutionary Search
G Require: Initial population of reward functions Po (r) 1==1, number of generations G. number of pairs M, candidates per pair N, selection size K for generation g = i,.... G do
(,(r)))) from current population Pg-1 2: Randomly sample M pairs {(r]))) (rpa)))
3: # Crossover Generation
4: for each pair 1,...,.u MM do
5: Use LLM to generate candidate reward functions (r-mm,a) /NV m==1
6: end for
7:
# Fitness Evaluation
8:
for each candidate (m, n²) do
(m.n) Run Algorithm with reward assignment function r
9:
(m,n) and Wasserstein distance D (m.)) Obtain policy Te,
10:
wasserstein.
11:
end for
12:
# Selection
13:
Select top K candidates {r}} with least divergence from the union of the current population
Pg-1 and generated crossovers (r (m,,)) }:
Pg + fri]k
end for
15: return Best reward function(s) rÿ and corresponding learned policies Tेr
20

---

Published as a conference paper at ICLR 2026
C
PROMPT STRATEGY
C.I TEMPLATE
We use the following prompt to instruct the LLM to generate crossover RA functions:
Prompt for Crossover Generation
Role: Al Reseurch Assistant (Imitation Learning)
Overall Objective: Collahorate to discover novel reward functions for Adversarial Imiitation I Learning (ATL) that improve training stability final policy performance. Performance measured by performance scure (higher betler).
Background: Adversarial Imitation I Learning Setting
have a policy ı and expert transitions (s, a] stored dataset Dg. typical learning loop involves:
Sampling transitions (s, a) into dataset D= using the current policy
Training a discrimiiator D(s, a) to distinguish between expert transitions (Dz) and policy transitions (Dx) using
standard binary entropy loss:
L - () (Do a))] at [Bog(1 D(e, a))]
The discriminator's ourput logits, a). approximate log-density
I(s, a) log 0 (s,a) (*)
Policy transitions (s) in D, are assigned rewards based on these logits using reward function r(s, a)
[(lls, aj). Examples include:
GAlL: r(s, a) = - log(1 D(s. ))) softplua(((s.e) e)) (Smooth rectifier. near O for negative logits,
linear positive).
AIRL: r(s, a) = log D(s.a) a) log(1 D(s,a)) [(s, a) (Linear everywhere).
FAARL: r(,, a) [(s, a) exp([(s a)) (Rises from to 1/e at = -1, then drops sharply).
log D(s, a) -sotppuus( (=(plluss((s.... a)) (Lineur for negalive logits, for positive). The policy x is updated using reinforcement learning (e.g. PPO, SAC) with these calculated rewards.
LOGD:
1-5 repeated.
Your Task in This Interaction:
reward lincctions, fi fo (delined Your goal is to propose a "new* function (not the same as GAT,, AIRI,, FAIRI, LOGD), fas that aims to perform hetter (higher scure).
presented
logits along their observed performance.
Instructions:
1. Analyze f and f2:
Consider mathematical shapes and properties (e.g.. monotonicity, bounds, smoothness).
Consider their behavior when the logits are near zero, positive, and negative. What signal do they provide? provided performance data. Why might Design fa reward. In (logits):
these properties
performed better/wurse?
Based on your analysis, propose "new* function fa-
Aim for diversity: Propose of novel functions and variations the provided examples.
Implementation Requirements:
Input: logits (a JAX array).
Output: zexard (a JAX aray of the same shupe).
Language:
Function Name: reward.In.
Clarity: Ensure the code is cleam, well commented (if neceswrry), and easily extractable. Include necessary jap,
imports (: Çax paax. nurpy
etc.).
Enclose in Code Block: Use a code block with the language python.
Jittable:
function jittable
Response Format:
import jax.numpy a5 JRP
from jax import an Uncomment or other importa 1" needed
def seward_fn
*en
[Rrief description of the function's togic/intent nen
# (Your implementation here] reward seturn zeward
Pair of Reward Functions:
Function I:
[Code]
Score: [scoro]
Function
[Code]
Score: [score]
21

---

Published as a conference paper at ICLR 2026
C.2
GENERATED EXAMPLE:
Below an example RA function generated by the LLM using the previously described prompt.
Parent
Inport jax
import jax.numpy jap import jax.nn
def reward_fnllagittn
Softplua function, clipped
(e.g., 5.0),
Prevents potentially destabilizing large positive zewards. (in (logita), (LIP_VALUE)
CLIF_ 5.0
jup.minimum wwt.:/a......pppp.a...oo wwwp:.nnnoottllsaaoo.. (logits), CLIP VALUE)
Parent
inport jax. as jnp
reward_ (logits): reward -logitas oxp. (logits) return reward
Score: -15.0
Generated Crossover
import nuppy import jox.nn
def feward_fn[logize
Combinea smooth non-linear shaping with symmetric bounds:
Uses tenh scaled by softplus to keep rewards smooth and bounded. negative logits, zewards
but smocthly increase toward small positive values, avoiding harsh penaltios. positive logits,
saturate mitigoting exploding pasitive rewards that couse instability. Intuttion:
tanh,
This function gently encouzages the policy Lowards regiona where logits positive,
but heavily penalizing large values is avoided by the tanh saturation. Provides smoother gradients
by multiplying softplus (logits) which behaves
ke Rat0.
***
softplus ansures smooth soft-reatific for logics, avoiding zero-gradient dead zones. "0 for large regacive, Linear for large positive
89 aax.mnooftppus
tanh bounds output between -1 and 1, adding mild negative zewards for strongly negative logits Scale tanh (logisl) to control sharpneas:
scaled_Lanh Lap. (legits
Combine signals multigiicatively:
reward scaled_tanh
reward
D HYPERPARAMETERS
The hyperparameters used across different components are summarized below. Table 4 contains the parameters for the LLM evolution process. The PPO hyperparameters used to train both the expert and the AIL policy are listed in Table 5, and the discriminator-spennee settings are provided in Table 6. The hyperparameters for the A2C algorithm reported in Table 7.
22

---

Published as a conference paper at ICLR 2026
Table 4: LLM Evolution hyperparameters
Parameter
Value
LLM Model
GPT-4.1-min
Number of Generations
Number of Pairs
20
Candidates per Pair
Selection Size
10
Number of evaluation seeds
16
Table 5: PPO hyperparameters for MinAtar and Brax
Hyperparameter
MinAtar
Brax
Number of Environments
64
2048
Number of Env Steps
128
10
Total Timesteps
x 10
10
Number of Minibatches
32
Number of Epochs
Discountt Factor
0.99
0.99
GAE
0.95
0.95
PPO Clip
0.2
0.2
Value Function Coefficient
0.5
0.5
Entropy Coefficient
0.01
0
Max Gradient Norm
0.5
0.5
Layer Width
64
256
Number of Hidden Layers
Activation
relu
relu
LR
0.005
0.0003
Anneal LR
linear
Optimizer
adam
adam
E INDIVIDUAL TRAINING CURVES
Figure presents the training curves for DAIL, and baseline methods. DAIL consistently outper- forms the baselines across all evaluated environments, with the exception of AIRL outperforming DAIL on Reacher. Additionally, DAIL. exhibits faster convergence even when final returns are comparable, underscoring the training stability introduced by its meta-learned reward assignment function.
F DISCOVERED REWARD ASSIGNMENT FUNCTIONS
The top five reward assignment functions discovered are shown in Table 8. While the resulting functions are complex, they are partly composed of primitives found the base population, such as log and e along with ones such as |x], min(x, y) and max(x, y).
G RUNTIME AND COMPUTE USED
Our experiments were conducted using a mix GPUs available on our compute cluster, including NVIDIA L40, A100, 3090, and H100 NVL. Each AlL evaluation involved training 16 agents in parallel, distributed equally across two GPUs. The wall-clock time for each evaluation reported in Table 9.
H COMPARISON WITH OPEN-ES
23

---

Published as a conference paper at ICLR 2026
SpaceInvaders
Breakout
Normalized Retur
1.0
Normalized Return
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
1M
3M
5M
7M
10M
1M
3M
5M
7M
10M
Steps
Steps
Asterix
hopper
Normalized Return
1.0
Normalized Return
1.5
0.8
0.6
1.0
0.4
0.5
0.2
0.0
0.0
1M
3M
5M
7M
10M
1M
13M
25M
37M
50M
Steps
Steps
walker2d
ant
Normalized Return
1.0
Normalized Return
1.0
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
1M
13M
25M
37M
50M
1M
13M
25M
37M
50M
Steps
Steps
halfcheetah
reacher
Normalized Return
1.0
Normalized Return
1.2
0.8
1.0
0.8
0.6
0.6
0.4
0.4
0.2
0.2
0.0
0.0
1M
13M
25M Steps
37M
50M
1M
13M
25M Steps
50M
DAIL
GAIL
GAIL-h
AIRL
FAIRL
Expert
Figure 9: Mean normalized returns across all evaluated environments. DAIL consistently outper- forms baseline methods, with the exception of AIRL Reacher.
24

---

Published as a conference paper at ICLR 2026
Table 6: Discriminator hyperparameters for Brax and Minatar,
Hyperparameter
MinAtar
Brax
Layer Width
64
128
Number of Hidden Layers
1
Activation
relu
relu
Learning Rate (LR)
0.0003
0.0003
Gradient Penalty Weight
0.1
1.0
Number of Epochs
1
1
Number of Minibatches
8
32
Activation
relu adam
relu adam
Optimizer
Table 7: A2C hyperparameters for MinAtar SpaceInvaders
Hyperparameter
Value
Number of Environments
64
Number of Env Steps
16
Total Timesteps
x 10
Number of Minibatches
Discount Factor
0.99
GAE À
0.95
Value Function Coefficient
5.0
Entropy Coefficient
0.01
Max Gradient Norm
10.0
Layer Width
64
Number of Hidden Layers
Activation
relu
LR
0.005
Anneal LR
linear
Optimizer
adam
We evaluate the contribution of the LLM in guiding the evolutionary search by comparing widely used baseline for optimizing the outer-loop objective (Eq.3): OpenAl Evolution Strategies (ES) Salimans et al. (2017). OpenAI- ES is a black-box, gradient-free optimization algorithm known for strong performance on similar classes of problems Goldie et al. (2024); Sapora et al. (2024). For fair comparison, both methods are given the same computational budget (200 inner-loop evaluations). From Ta- ble 10, we observe that OPEN-ES underper- forms DAIL both on the training environment (Spacelnvaders) and on the test environments overall, consistent with the findings reported by Goldie al. (2025). Figure 10 visualizes the discovered RA function, which appears irregular and non-smooth, likely contributing to its inferior performance.
15
10
r|s,a
55
-3
-2
1
log
(viS.
Figure 10: RA function generated by OpenES
I COMPARISON WITH ADDITIONAL DIVERGENCE BASED METHODS
We evaluate RA functions derived from alternative f-divergences, as summarized Table 3, which has not been explored in prior work Ghasemipour et al. (2020); Orsini et al. (2021). We also compare
25

---

Published as a conference paper at ICLR 2026
Table 8: Top reward assignment functions discovered generated after evolution, Each function is expressed in terms of discriminator logits I.
Environment
Reward Assignment Functions
Spacelnvaders-Run 1 1. (()) 0.5 1155.(aann (tanh(!) +1)
softplus (16((1-0.8))
0.5 0.81
1 -08
1.5
-0.8 < < < 0.8 ))
2. min
1.5, max
0,
0.5 0
softplus
(.5(1-0.8)
0.5 + 0.8 0.8 +
12 0.8
1.5
3. softplus(I)
(1.5/) + 0.5 gelu(()
(31) 0.5 (tanh(l) + 1)
++
5. 0.5
1+ [2] + +1) (31))
Spacelnvaders-Rui 2
1. clip(0.5
(tanh (1) + 1) softplus(() 0.1 softplus(-1), 0, 00)
(1 (1- tanh?
1-3
2. softplus(()
1.5
3. (31) softplus(1)
4. a(3(1 - 1)) ([5(255-11))
clip (11/11 + 0.2
1
3
+ 0.3) exp( 0.2 max(1, 0)). 00
1+10]
Breakout
1. 0.5 (tanh(1.5- +))+))+)) +1
2. softplus(1) 1 (l) +1
2
3. (tanh(I) + +1) (I))
tanh(2.0- + 0.3 softplus(I) + 0.5
2.0
5. a2.00 1)) 0.7 softplus() + a2.00-1) 1) (softplus(() + cl.p -1,,2)
Table 9: All reported runtimes correspond to wall-clock time measured while training 16 agents in parallel on 2 H100 NVL GPUs.
Environment
Wallelock time
Ant
346.53
HalfChectah
1034.86
Hopper
585,97
Walker2d
701.22
Reacher
290.54
Breakout
55,07
Asterix
97.34
SpaceInvaders
52.57
DAIL to Wasserstein-GAlI. (WGAIL), which corresponds to replacing the discriminator objective
26

---

Published as a conference paper at ICLR 2026
Table 10: Comparison between DAIL and OPEN-ES across MinAtar and Brax environments. Re- ported values denote mean + standard error over evaluation runs. Note that DAIL and OPEN-ES were meta-trained on MinAtar SpaceInvaders.
Environment
DAIL
OPEN-ES
SpaceInvaders'
0.90 + 0.00
0.73 0.05
Asterix
0.66 +
1.27 =
Breakout
1.01 1: 0.00
0,38 : 0.10
HaltCheetah
0.90 1 0.00
0.62 = 0.02
Walker2d
0.95 1 0,00
0.73 0,02
Hopper
1.65 0.01
1.09 0.03
Reacher
0.63 + 0.01
0.83 =
Ant
0.85 + 0.02
0.51 = 0.03
Table 11: Comparison with additional divergence-based AIL methods across MinAtar and Brax environments. Reported values denote mean + standard error.
Environment
DAIL
Pearson
Sq- Sq-Hellinger
TV
WGAIL
Asterix
0.66
0.03
-0.03
0.00
-0.02 0.00
-0.01 0.00
0,52 +: 0.04
Breakout
1.01 = 0.00
-0.01 = 0.00
+0.01 = 0.00
+0.01 = 0.00
0.91 + 0.00
Ant
0.85 0,02
0.57 = 0.02
0.82 + 0.03
0.19 i 0.05
0.80 + 0.00
HalfCheetah
0.90
0.49
0.86 001
-0.01
0.87 0.00
Hopper
1.65
0,90
1.49
1.59
1.50
Reacher
0.63 = 0.01
0.88 + 0.05
0,94 + 0.01
0,86 + 0.07
0.42 + 0.01
Walker2d
0.95 0.00
0.54 + 0.05
0.98 + 0.00
0.91 + 0.01
0.89 + 0.01
(36) with:
D) arg
max S1 ((s,a)~gg [D(s,a)] B(saa)p a)
(38)
Note that we already employ gradient penalty as a regularizer, enforcing the discriminator to be 1- Lipschitz. After hyperparameter sweep over the discriminator learning rate, find that 3 10-3 performs best for WGAIL. From Table 11, we observe that DAIL outperforms the divergence-basee AIL baselines on 5 out of 7 test environments.
J COMPARISON WITH
METHODS
We additionally compare DAIL with non-adversarial methods, namely Behavior Cloning (BC; Pomerleau (1988)) and IQ-Learn (Garg et al., 2021), a state-of-the-arr non-adversarial algorithm. For IQ-Learn, we use the official implementation provided by Garg al. (2021). Results are re- ported in Table 12. DAIL consistently outperforms both BC and IQ-Learn. Interestingly. IQ-Learn fails to surpass BC in 3 of the tested environments, phenomenon also observed in prior work (Lai al., 2024; Jain al., 2024), highlighting stability challenges inherent in con-adversarial imitation learning methods as well.
Table 12: Comparison with von-arversarill imitation learning algorithms. Reported values are mean returns + 95% confidence intervals across 4 random seeds for each environment.
Method
Ant
HalfCheetah
Hopper
Walker2d
BC
0.23 1: 0.00
0.09 1: 0.02
0.23 009
0.04 :: 001
IQ-learn
-0.32 +: 0.0
-0.02 1: 0.00
0.02 006
0.03 : 001
DAIL
0.88 +: 0.06
0.90 + 0.01
1.72 004
0.97 + 001
27

---

Published as a conference paper at ICLR 2026
K
LLM USAGE
LLMss have been used in the writing of this paper, primarily to refine the quality of the text through prompt-based polishing of author-written drafts.
28


Reviews:
[{"review_id": "SsKWiGof66", "reviewer": "Reviewer_5SFa", "summary": "This paper focuses on discovering useful reward functions for adversarial imitation learning, with the goal of improving downstream performance of agents beyond that of hand-designed IL baselines that rely on f-divergence minimization. The paper proposes DAIL, the first meta-learned AIL algorithm, which outperforms standard AIL algorithms on a variety of tasks in continuous control (Brax) and in Atari (MinAtar).\n\nDAIL formulates the meta-learning problem as minimizing the Wasserstein distance between the expert distribution induced by the dataset and the optimal policy with respect to the meta-learned reward function. The meta-learning approach is evolutionary search, which is simulated by an LLM that takes in two \"parent\" reward functions and aims to produce a \"child\" reward function that combines desired properties of both parents, making it more suitable for downstream IRL. Once a said reward function is given, standard AIL proceeds, and the Wasserstein distance between the final learned policy and the expert distribution is estimated, which is signal for the meta-learner.", "questions": "I don't have any major questions right now -- I think the paper is pretty clear and I think I understood the main points reasonably clearly.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is very cleanly written, and even as a person who doesn't have much of a meta-learning or metagradient RL background, I was easily able to understand the problem and semi-verify that the approach was reasonable. The approach is also fairly novel -- as the paper claims, and as I have verified, there has been no prior work that has attempted to meta-learn a reward function via standard meta-gradient or evolutionary search approaches in AIL.\n\nThe results in the paper are also quite strong, improving upon standard baselines such as GAIL and AIRL somewhat considerably based on the probability of improvement plot in Figure 5, meaning that the meta-learning approach is yielding strong reward functions in practice.", "weaknesses": "For the scope of this paper, I don't really know if this should count as a weakness, but I would have liked to see some additional baseline comparison to other AIL algorithms beyond AIRL and GAIL. In particular, there are many possible f-divergences or IPMs that could be minimized to obtain a strong adversarial IL algorithm in practice, including the Wasserstein distance (e.g. see Wasserstein GAN [1]). This is in particular a useful baseline because the meta-learning objective is aiming to minimize the Wasserstein distance between the expert policy and the learned policy, which a Wasserstein AIL objective building on top of the Wasserstein GAN would in theory do. This has also been done in other work, such as the hybrid IRL paper [2].\n\nIn addition, it may have been a good idea in my opinion to see how the meta-learned reward function translated to much more complicated tasks than those of Brax and MinAtar. This was done in the Atari suite by Oh et. al. [3] when they discovered RL update rules, so maybe from an experimental standpoint it would be a good idea here. However, this is not a huge weakness in my eyes as Brax and other MinAtar environments are already much different than MinAtar SpaceInvaders, which was used for meta-training. Additionally, maybe some work could've been done focusing on if meta-learning the reward function across more environments would've led to a more robust RL algorithm.\n\n[1]: Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville: Improved Training of Wasserstein GANs (NIPS 2017)\n[2] Juntao Ren, Gokul Swamy, Zhiwei Steven Wu, J. Andrew Bagnell, and Sanjiban Choudhury: Hybrid Inverse Reinforcement Learning (ICML 2024)\n[3] Junhyuk Oh, Matteo Hessel, Wojciech M. Czarnecki, Zhongwen Xu, Hado van Hasselt, Satinder Singh, and David Silver: Discovering Reinforcement Learning Algorithms (NeurIPS 2020).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GNJ7eUeLXx", "reviewer": "Reviewer_6x5G", "summary": "This paper addresses the long-overlooked reward assignment (RA) function design in Adversarial Imitation Learning (AIL) by proposing an LLM-guided evolutionary search framework. The authors formalize RA function discovery as a meta-learning problem optimizing Wasserstein distance, conducting the search on MinAtar SpaceInvaders to discover DAIL's RA function: r(x)=0.5·sigmoid(x)·[tanh(x)+1]. Experiments demonstrate that DAIL outperforms GAIL, AIRL, FAIRL, and other baselines across multiple unseen environments in both Brax and MinAtar suites. The authors further analyze policy entropy and log-ratio distributions to elucidate the mechanisms behind DAIL's enhanced training stability.", "questions": "1. If the search were conducted independently on Brax environments such as Ant or Walker, would the resulting RA functions be similar in form to the DAIL function discovered on SpaceInvaders? How large would the performance differences be? This experiment is crucial for validating the method's true generalization capability. Relatedly, are the performance differences among the top-5 functions shown in Table 6 statistically significant? \n\n2. Although DAIL does not strictly correspond to any f-divergence, could the authors provide intuitive analysis or informal reasoning about convergence under relaxed assumptions (e.g., small learning rates, approximately optimal discriminators, bounded state spaces)? Even without proving convergence, could the authors explicitly discuss the types of environments or conditions where DAIL might fail (e.g., highly non-stationary environments, tasks with extremely sparse rewards)? \n\n3. If multiple independent searches were conducted (using different random seeds to initialize the LLM or population), how high would the overlap be among the top-K functions? This would reflect the stability of the search process. Additionally, could the authors report the complete computational cost (total GPU hours, rather than the per-evaluation time shown in Table 7)? Given 200 candidate functions, 10 generations of evolution, and each evaluation requiring training 16 agents to convergence, the total cost could be substantial. Clarifying this is important for assessing the method's practicality. \n\n4. Why does IQ-Learn's performance in Table 8 significantly underperform its original paper's reports (even failing to beat simple BC in multiple environments)? Has the implementation correctness been verified? How were the hyperparameters tuned? This anomalous result weakens the credibility of comparisons with non-adversarial methods. Additionally, since the paper uses Wasserstein distance as both the fitness function and evaluation metric, could the authors specify what cost matrix is used (Euclidean distance? or other?)", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. The AIL literature has predominantly focused on discriminator stabilization (Wang et al. 2024; Luo et al. 2024), while RA function design has been consistently overlooked. This paper is the first to systematically treat RA functions as an independent research subject, filling a significant void. More importantly, the authors shift the RA design paradigm from theory-driven (deriving from f-divergences) to data-driven (directly optimizing performance metrics), offering a fresh methodological perspective for AIL algorithm design. The discovered DAIL function is remarkably concise (one line of code) and easily integrable into existing AIL frameworks, demonstrating high practical value.\n\n2. The choice of Wasserstein distance as the fitness function is natural, given its theoretical robustness in measuring distributional discrepancies. The LLM-guided evolutionary search organically combines the interpretability of code representation with search efficiency, providing more targeted exploration compared to pure random search. The initial population includes classical methods (GAIL, AIRL, FAIRL, GAIL-heuristic), ensuring high-quality starting points—a design choice that reflects the authors' deep domain understanding. Overall, from problem formalization to search strategy selection, the design demonstrates sound intuition.\n\n3.  The training dynamics analysis in Figure 6, examining policy entropy and log-ratio distributions, clearly reveals why DAIL outperforms baselines: DAIL saturates to near-zero in the log-ratio < -1.8 region, effectively filtering noisy signals corresponding to random or low-quality behaviors, thereby achieving more stable training. This mechanistic explanation is not only data-supported (policy entropy converges faster to levels approaching PPO with true rewards) but also holds consistently across multiple environments (though space constraints limit presentation to selected examples), strengthening the credibility of the conclusions. The ablation study in Figure 7 further confirms the necessity of combining sigmoid and tanh, showing that using either function alone cannot match the combined effect.", "weaknesses": "1. The lack of theoretical guarantees is the most serious issue in this work. DAIL's RA function does not correspond to any known f-divergence, and thus lacks convergence guarantees. While the paper acknowledges this limitation (Lines 247-251), it merely justifies the approach with \"empirical feedback\" and \"robust generalization\" without providing any alternative theoretical analysis. In contrast, GAIL, AIRL, and FAIRL all have formal proofs of convergence to ρ_π = ρ_E (Ghasemipour et al. 2020). This deficiency means we cannot predict under what conditions DAIL will be effective or when it might fail, nor can it guide future RA function design. For an algorithmic paper, relying entirely on empirical results without theoretical foundation is a critical weakness. Even if rigorous convergence proofs are intractable, the authors should at least provide intuitive analysis under relaxed assumptions (e.g., bounded MDPs, Lipschitz continuous policy classes, approximately optimal discriminators) or explicitly discuss potential failure modes and applicability boundaries.\n\n2. The evidence for generalization is fundamentally insufficient. The paper's core contribution claims to have \"discovered a generalizable RA function,\" but the evidence supporting this claim has serious problems. According to Section 6.1 (Lines 347-353), the authors conduct the search on only one environment—MinAtar SpaceInvaders—and then test on the other 7 environments (5 Brax + 2 other MinAtar). This raises two critical concerns: First, SpaceInvaders is an Atari game with discrete action space, while Brax environments involve continuous control tasks with entirely different dynamics. That an RA function discovered on SpaceInvaders performs well on Brax could be a coincidence (SpaceInvaders happens to be a \"good\" meta-training environment) rather than evidence of DAIL's true cross-domain generalization capability. Second, different task types (e.g., sparse vs. dense reward environments, or varying state space dimensions) may require RA functions with different characteristics, yet the paper does not explore this at all.\n\n3. The missing critical experiments make the generalization claims unconvincing: We would expect to see searches conducted independently on at least 2-3 different environments (e.g., search on Ant then test on Walker/HalfCheetah; search on Breakout then test on Asterix/SpaceInvaders), followed by comparison of whether the discovered RA functions are consistent. Only if multiple independent searches converge to similar functional forms would that constitute genuine evidence of generalization. The current experimental design at best demonstrates that \"SpaceInvaders is a good meta-training task\" but cannot establish \"DAIL's generalization capability.\" The performance improvements shown in Figure 4 may merely reflect the peculiarities of SpaceInvaders rather than the method's universality.\n\n4. Several writing and presentation details could be improved. Upon verification, the paper is largely consistent in notation usage (primarily using rfr_f\nrf​ and ℓ\\ell\nℓ), but the statement \"offers less flexibility\" on Line 85 is overly vague and requires more specific explanation of which aspects of non-adversarial methods lack flexibility. The title of Appendix A.2 presents potential confusion: since the paper uniformly uses the form Df(ρE∣∣ρπ)D_f(\\rho_E || \\rho_\\pi)\nDf​(ρE​∣∣ρπ​), whether FAIRL should be labeled as \"Forward KL viewed as reverse\" or simply \"Forward KL\" depends on the reader's interpretive perspective. The authors should clarify this explicitly in the title or main text to avoid misunderstanding.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "x45KreeYPu", "reviewer": "Reviewer_JUMG", "summary": "This paper proposes Discovered Adversarial Imitation Learning (DAIL), a meta-learning framework that aims to automatically discover reward assignment (RA) functions for adversarial imitation learning. \nDAIL employs an LLM-guided evolutionary search to generate and evolve candidate RA functions expressed as executable code, optimizing them based on the Wasserstein distance between expert and policy occupancy measures. \nExperiments on Brax and Minatar benchmarks show that DAIL achieves higher normalized returns and better alignment with expert behavior than existing AIL baselines. \nThe paper claims to offer the first meta-learned AIL algorithm and provides empirical analyses linking the discovered reward structure to more stable adversarial training dynamics.", "questions": "1. Why does the method use the Wasserstein distance as the evaluation objective while still training within the f-divergence–based AIL framework? If Wasserstein distance provides smoother and more stable gradients, why not directly optimize it instead of meta-learning new reward assignment functions?\n\n2.How does the discovered reward function r theoretically relate to any valid divergence or distance measure?\n\n3. Does the use of Wasserstein distance for evaluating fitness bias the search toward functions that implicitly approximate it?\n\n4. How does DAIL compare in stability and performance to Wasserstein-based imitation methods [5, 6]?\n\n5. Can the instability in AIL truly be mitigated by changing the RA, or does it fundamentally arise from the adversarial optimization dynamics?\n\n6. Would replacing the adversarial discriminator with a 1-Lipschitz critic (as in WGANs) render the entire reward-assignment discovery unnecessary?\n\n[5] Zhang, Ming, et al. \"Wasserstein distance guided adversarial imitation learning with reward shape exploration.\" 2020 IEEE 9th Data Driven Control and Learning Systems Conference (DDCLS). IEEE, 2020.\n[6] Dadashi, Robert, et al. \"Primal wasserstein imitation learning.\" arXiv preprint arXiv:2006.04678 (2020).", "limitations": "", "rating": 2, "confidence": 5, "soundness": 1, "presentation": 3, "contribution": 1, "strengths": "1. *Novelty.* The paper presents a moderately original idea by framing RA discovery in AIL as a meta-learning problem and leveraging LLMs for evolutionary search over reward function space. Its originality stems primarily from the methodological integration of LLM-guided symbolic evolution into the imitation learning pipeline, rather than from a fundamentally new theoretical insight about AIL itself. \n\n2. *Quality.* The paper is technically sound in its experimental setup and provides comprehensive empirical validation on both Brax and Minatar benchmarks, demonstrating consistent improvements in performance over standard baselines. \n\n3. *Clarity.* The paper is clear and easy to follow. \n\n4. *Significance.* The work contributes an interesting demonstration of how LLMs can aid in discovering algorithmic components in AIL. However, its broader impact on the IL/RL community is modest, as the main performance gains stem from empirical search rather than conceptual advances, and the underlying instability of AIL that is driven by the adversarial optimization process, remains unresolved.", "weaknesses": "The main weakness of this paper lies in its limited theoretical novelty and misattribution of the source of instability in AIL. \n\n1. The idea of studying how RA functions affect learning dynamics is not new. Prior works [1,2,3] have already analyzed and benchmarked multiple RA designs under the f-divergence framework. \n\n2. The proposed contribution of using an LLM-guided evolutionary search to discover new RA functions only offers a methodological novelty rather than a theoretical or algorithmic one. However, this approach does not address the primary cause of instability in AIL, which stems from the adversarial f-divergence min–max optimization dynamics rather than the form of the f-divergence [4]. As a result, the proposed method may yield empirically smoother rewards but cannot fundamentally stabilize training.\n\n3. The discovered RA functions violate theoretical constraints that ensure valid divergence minimization, which undermines interpretability and convergence guarantees. The paper could be significantly strengthened by incorporating these mathematical constraints into the search process (e.g., by enforcing convexity) or by analyzing the discovered functions’ theoretical properties relative to valid divergences.\n\n4. Empirically, while the experiments on Brax and Minatar demonstrate some improvements, the evaluation is limited to relatively simple continuous-control and grid-world tasks. The generality and robustness of DAIL would be more convincing if tested on high-dimensional or visual IL environments (e.g., Atari) and under different discriminator architectures or regularization strategies. \n\n5. The paper attributes stability improvements solely to the discovered reward function, but it lacks ablation studies isolating the effect of LLM-guided search versus simpler baselines (e.g., random search, symbolic regression, or gradient-based meta-learning). Including such comparisons would clarify whether the LLM’s involvement meaningfully contributes beyond automated function optimization.\n\n[1] Ke, Liyiming, et al. \"Imitation learning as f-divergence minimization.\" International workshop on the algorithmic foundations of robotics. Cham: Springer International Publishing, 2020.\n[2] Ghasemipour, Seyed Kamyar Seyed, Richard Zemel, and Shixiang Gu. \"A divergence minimization perspective on imitation learning methods.\" Conference on robot learning. PMLR, 2020.\n[3] Zhang, Xin, et al. \"f-gail: Learning f-divergence for generative adversarial imitation learning.\" Advances in neural information processing systems 33 (2020): 12805-12815.\n[4] Arjovsky, Martin, and Léon Bottou. \"Towards principled methods for training generative adversarial networks.\" arXiv preprint arXiv:1701.04862 (2017).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dg65tpq0Et", "reviewer": "Reviewer_GF4H", "summary": "This paper proposes Discovered Adversarial Imitation Learning (DAIL) that stabilizes the training process by using an LLM-guided framework to explore the Reward Assignment (RA) functions.", "questions": "- Some of the text in Figure 2 is too small to read, which affects the clarity of the visualization. Improving figure readability would strengthen the presentation.\n- There seems to be a contradiction regarding the reward assignment (RA) function: in Section 5.2, it appears to change dynamically, while the future work section mentions that the RA function in DAIL is static. Could the authors clarify whether the RA function is adaptive during training or fixed throughout?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This paper presents a search procedure that illustrate how the LLM-guided framework searches for the reward assignment functions\n- Empirical experiments on MuJoCo and Minatar tasks have show the effectiveness of DAIL", "weaknesses": "- The paper primarily builds upon existing reward assignment strategies, and thus the level of novelty is limited.\n- While experiments support the proposed idea, the paper lacks rigorous theoretical analysis to clarify to what extent DAIL improves stability or under what conditions the improvements are guaranteed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)