GRACE: A Language Model Framework for Explainable Inverse Reinforcement Learning
Authors: Silvia Sapora
R Devon Hjelm
Omar Attia
Alexander T Toshev
Bogdan Mazoure
Venue: ICLR 2026 Poster (2026)
Keywords: inverse reinforcement learning
large language models
evolution
Primary area: reinforcement learning
TL;DR: We introduce GRACE, a framework that uses code-generating language models within an evolutionary search to learn an interpretable reward function as executable Python code directly from expert demonstrations

Abstract: Inverse Reinforcement Learning aims to recover reward models from expert demonstrations, but traditional methods yield black-box models that are difficult to interpret and debug. In this work, we introduce GRACE (**G**enerating **R**ewards **A**s **C**od**E**), a method for using Large Language Models within an evolutionary search to reverse-engineer an interpretable, code-based reward function directly from expert trajectories. The resulting reward function is executable code that can be inspected and verified. We empirically validate GRACE on the MuJoCo, BabyAI and AndroidWorld benchmarks, where it efficiently learns highly accurate rewards, even in complex, multi-task settings. Further, we demonstrate that the resulting reward leads to strong policies, compared to both competitive Imitation Learning and online RL approaches with ground-truth rewards. Finally, we show that GRACE is able to build complex reward APIs in multi-task setups.

Published as a conference paper at ICLR 2026
GRACE: A LANGUAGE MODEL FRAMEWORK FOR EXPLAINABLE INVERSE REINFORCEMENT LEARNING
Silvia Sapora
Devon Hjelm
Alexander Toshev
Omar Attia
Bogdan Mazoure
ABSTRACT
Inverse Reinforcement Learning aims to recover reward models from expert demonstrations, but traditional methods yield black-box models that are difficult interpret and debug. In this work, we introduce GRACE (Generating Rewards As CodE), method for using Large Language Models within an evolutionary search reverse-engineer an interpretable, code-based reward function directly from expert trajectories. The resulting reward function is executable code that can be inspected and verified. We empirically validate GRACE on the MuJoCo, BabyAl AI and AndroidWorld benchmarks, where it efficiently learns highly accu- rate rewards, even in complex, multi-task settings. Further, we demonstrate that the resulting reward leads to strong policies, compared to both competitive Imita- tion Learning and online RL approaches with ground truth rewards. Finally, we show that GRACE is able to build complex reward APIs in multi task setups.
INTRODUCTION
The performance of modern Reinforcement Learning (RL) agents is determined by, among other factors, the quality of their reward function. Traditionally, reward functions are defined manually as part of the problem specification. In many real-world settings, however, environments are readily available while reward functions are absent and must be specified. Manually designing rewards is often impractical, arror-prone, and does not scale, particularly in contemporary multi-task RL. scenarios (Wilson et al., 2007: Teh et al., 2017 Parisotto et al., 2016).
A natural alternative is to automate reward specification by learning reward model from data. The dominant paradigm here is Inverse Reinforcement Learning (IRL), which attempts infer a reward model from observations of expert behavior (Ng & Russell 2000 Christiano et al. 2017 Ziebart ct al. 2008). In the era of Deep RL, approaches such as AIRL (Fu et al.] 2018) represent rewards with deep neural networks. While effective, these reward functions are typically opaque black boxes, making them difficult to interpret or verify (Moinar) 2020). Moreover, IRL methods often require substantial amounts of data and can lead inaccurate rewards (Sapora al. 2024).
Recently, expressing reward models through code has emerged as promising approach Venuto al. 2024a Ma et al. 2023). Code is particularly well-suited representation, reward functions
are often far simpler to express than the complex policies that maximize them (Ng & Russell 2000)
Cook [971 Godel 1956). These approaches leverage code-generating Large Language Models
(LLMs) and human provided task descriptions or goal states to generate reward programs (Venuto al. 2024a). Subsequently, the generated rewards are verified (Venuto et al.) 2024a) or improved using the performance of a trained policy as feedback (Ma et al. 2023). However, this prior work has not investigated whether it is possible to recover a reward function purely from human demon- strations in an IRL-style setting, without utilizing any explicit task description or domain-specific design assumptions.
In this work, we address the question of how to efficiently infer rewards-as-code from expert demon- strations using LLMs. We propose an optimization procedure inspired by evolutionary search (Gold- berg 1989b Eiben & Smith 2003b Salimans et al. 2017a| Comera-Paredae et al.) 2024a Novikov et al. 2025b). in which code LLMs iteratively introspect over demonstrations to generate and refine as reward models. This perspective effectively revisits the IRL paradigm in the modern context of program synthesis with LLMss.
programs that
Our contributions are threefold:

---

Published as a conference paper at ICLR 2026
Sample-Efficient Reward Recovery: We demonstrate that code LLMs conditioned on expert demonstrations can produce highly accurate reward models that generalize well to held-out data. Crucially, GRACE is highly sample-efficient, recovering accurate rewards from relatively few demonstrations without requiring manual domain knowledge or human- in-the-loop guidance.
Well-Shaped Rewards: We show that rewards learned by GRACE enable the training of strong policies across diverse domains, including the procedural BabyAI, continuous con- trol in MuJoCo, and real-world device control in AndroidWorld. Empirical results indicate that GRACE matches outperforms established IRL baselines (e.g., GAIL) and online RL with ground-truth, sparse rewards.
Interpretability and Modularity: Unlike black-box neural networks, GRACE generates rewards as executable Python code, making them inherently interpretable and verifiable. Furthermore, in multi-task settings, our evolutionary search naturally emerges a modular library of reusable reward functions, enabling efficient generalization across tasks.
2 RELATED WORKS
LLMs for Rewards A common way to provide verification/rewarrd signals in automated fashion is to utilize Foundation Models. LLM-based feedback has been used directly by Zheng et al. (2023) to score solution to critique examples (Zankner et al. (024). Comparing multiple outputs in relative manner has been also explored by Wang et al. (2023). Note that such approaches use LLM in shot fashion with additional prompting and potential additional examples. Hence, they can
utilize only small number of demonstrations at best. In addition to zero shot LLM application, it is also common train reward models, either from human feedback (Ouyang al. 2022) from AI feedback (Klissarov et al. 2023 2024). However, such approaches require training a reward model that isn't interpretable and often times require a larger number of examples.
Code as Reward As LLMs have emerged with powerful program synthesis capabilities (Chen
al. 202T Austin et al. 2021 Li et al. 2023 Fried et al. 2022 Nijkamp al. 2022) research has
turned towards generating environments for training agents (Zala et al. 2024 Faldor et al. 2025)
for various domains and complexities.
When it comes to rewards in particular, code-based verifiers use a language model to generate ex- ecutable Python code based on a potentially private interface such as the environment's full state. Because early language models struggled to reliably generate syntactically correct code, the first code-based verifiers (Yu et al. 2023 Venuto al. 2024b) implemented iterative re-prompting and fault-tolerance strategies. More recent approaches focus on progressively improving syntactically correct yet suboptimal reward function, particularly by encouraging exploration (Romera- Paredes et al. 2024b Novikkv et al. 2025a). Other approaches such Zhou et al. (2023): Dainese et al.
2024
use search in conjunction with self-reflection (Madaan et al. 2023) to provide feedback.
work, EUREKA (Ma al. 2023) assumes access to the ground-truth reward signal to evaluate and evolve the generated code, while Reward-As-Code (Venuto et al. 2024b) depends on explicit task descriptions goal states and hand written pipeline. In contrast to these earlier works, our pipeline doesn't 't require any domain-specific adaptation and does not require the ground truth reward signal description of the task.
While closely related
Inverse Reinforcement Learning (IRL) Early approaches infer a reward function that makes the expert's policy optimal over all alternatives (Ng Russell 2000). Subsequent works focused on learning policies directly, without explicit reward recovery (Abbeel & Ngl 2004), while incor- porating entropy regularization (Ziebart al. 2008) or leveraging convex formulations (Ratliff) et al. (006). More recently. deep IRL methods such as Generative Adversarial Imitation Learn- ing (GAIL) (Ho & Ermon 2016) have framed the problem as a distribution matching task using adversarial training. Adversarial Inverse Reinforcement Learning (AIRL) (Fu et al., 2018) further improves upon this by recovering disentangled and transferable reward functions. While related to our formulation, our representation (code) and our optimization strategy (evolutionary search) are fundamentally different, as GRACE benefits from implicit regularization through its symbolic reward representation.
2

---

Published as a conference paper at ICLR 2026
3 BACKGROUND
Reinforcement Learning We consider a finite-horizon Markov Decision Process (MDP) (Put- erman 2014) parameterized by M = AST,,T,' r) where S. A are the state and action spaces, (SS) the transition operator, and r is a reward function. The agent's behavior is described by the policy т S - A (A). Starting from a set of initial states Sp C S, the agent takes the action a r(() at 8, receives a reward rss) and transitions into state s' ~ as).
The performance of the agent is measured with expected cumulative per-timestep rewards, referred to as return:
H
J( = By m. r(s))
(1)
where 7 are trajectory unrolls of horizon H of the policy in M. An optimal agent can be learned by maximizing Equation [T) via gradient descent with respect to the policy, also known as policy gradient (Sutton et al. 1999. Schulman et al. (017).
Inverse Reinforcement Learning the reward r is not specified, it can be learned from demon- strations of an expert policy Tg. In particular, the classical IRL objective learns reward whose optimal return is attained by the expert (Ng & Russell 2000 Syed & Schapire 2007):
max min r ((, R m
More recent IRL approaches learn a discriminator that distinguishes between expert and non-expert demonstrations (Ho & Ermon 2016; Swamy et al. 202T). The likelihood of the agent's data under the trained discriminator can be implicilly thought of reward. Modern approaches often frame this as divergence minimization problem, matching the state-action visitation distributions of the learner and the expert. These approaches utilize gradient based methods to optimize their objectives.
As an alternative for cases where the objective is not readily differentiable, gradient-free methods can be employed. One such method is evolutionary search, which maintains set of candidate solutions (called population) and applies variation operators to improve it (Sal-) imans et al. 2017b Liben & Smith, 2003a; Goldberg 1989a). These operators include mutation, where hypothesis is partially modified, and recombination, where two hypotheses combined to produce a new one. Each variation is evaluated using a fitness function, which measures the quality of a given hypothesis. Starting with an initial population, evolutionary search repeatedly applies these variation operators, replacing hypotheses with higher-fitness alternatives.
Evolutionary search
In this work, we focus inferring reward functions, represented as Python code, from set of demonstrations. While this setup is related to IRL, representing rewards as code prevents us from applying gradient-based methods commonly used in IRL. For this reason, we adopt evolutionary search as our optimization method.
4 METHOD
We propose GRACE, Generating Rewards As CodE, an interpretable IRL framework that gener- ates a reward function as executable Python code. Initially, an LLM analyzes expert and random trajectories (Phase 1) and generates a preliminary set of reward programs. This initial set of reward functions is then iteratively improved through evolutionary search, where the LLM mutates the code based on low-fitness examples (Phase 2). The best-performing reward function is used to train an RL agent. This agent then explores the environment, and the new trajectories it generates are used to expand the dataset, revealing new edge cases or failure modes (Phase 3). This loop continues, with the expanded dataset from Phase 3 being used in the next iteration of Phase and 2, progressively improving the reward function. Notably, prior to the initial analysis, the LLM can optionally per- form a data cleaning step to identify and remove irrelevant states if the expert trajectories are known to be noisy or suboptimal. The overall process is illustrated in Figure I and detailed below and in Al Algorithm]
3

---

Published as a conference paper at ICLR 2026
Initialisation Phase I
Reward Refinement- Phase
RL Training Phase 3
Nogathe
Fibs Negativen
Reward Pusction
Reward
Palso Piallives
<
LAM
</>
Evalestico
Renard
LLM
<
repeat convergeace
Fitness
Balt Reward Fanction
Reward
Naxathre
New
Phase 2
Initialised
Figure Overview of the GRACE framework. (a) The expert, negative and extra data (if any) is used to generate an initial set of possible reward functions. (h) The expert and negative states are used to mutate reward functions through an evolutionary procedure. The rewards are iteratively refined by feeding low-fitness examples to the reward. (c) An agent is trained with online RL using the converged reward; the data it sees during the training is added to D and used to further improve the reward.
Phase Initialisation The initial reward code generation by GRACE is based on set of demon- stration trajectories DD and a set of random trajectories DD The former is generated using an expert policy or human demonstrations depending on the concrete setup, while latter is produced by random policy. Note that with a slight abuse of notation we will use 'D to denote interchangeably set of trajectories as well the set of all states from these trajectorics.
The language model is prompted with random subset of D and, optionally, extra information available about the environment (e.g. Python code or tool signature), to produce two artifacts:
Initial rewards: The LLM generates initial set Rinit of reward functions. Each function Rinit is represented as Python code:
def reward (state) -> float:
<LLM produced code>
designed to assign high values expert states Se and low values negative ones Sn, This set of rewards is treated as the population in the subsequent evolution phase.
(Optional) Data cleaning: The LLM analyzes the states from expert demonstrations to identify the subset of expert states Se S * that solve the task these are positive samples. All remaining states S, = (D+ S.) are treated as negative samples.
Phase 2: Reward Refinement through Evolutionary Search GRACE uses Evolutionary Search to discover rewards that effectively distinguish between expert and negative states. This is achieved by muluting the current reward population R using code LLM and retaining rewards with high fitness.
The fitness f of reward function r measures how well this function assigns large values to expert states and small values to negative states, akin to what would be expected from a meaningful reward. To do this, adopt the standard AIRL (Fu al. 2018) loss ensure transferable rewards:
f(r) Se [log (s)) E,-s., [lo([1 D,(s))
(3)
where (s) is the discriminator parametrised by the reward function r. Following the AIRL for- mulation, the discriminator takes the form (s) eaptr(e))+aaas)
exp(r(s))
The mutation operator utilizes LLM to improve the current reward code and correct failures.
We prompt the LLM to generate a mutated reward based on the current source code and specific failure cases:
m(r) = LLLM(source()), context, prompt)
(4)
Crucially, the context provided to the LLM enables targeted debugging. It includes:
4

---

Published as a conference paper at ICLR 2026
Current Reward Source Code: The Python implementation of the "parent" reward fune- tion r.
Wrong Examples: A random sample of states with low fitness under reward
Reward Values: The output value r(8w) by the current function for each low-fitness state.
(Optional) Extra environment information: It is possible include, if available, extra information about the environment such as (partial) source code or a text description.
(Optional) Debug information: The LLM can define debug(s, D* D)) function. The output of this function can then included, in-context, each mutation prompt. This allows the model to print custom intermediate variables, logic checks during execution or aggregate information over the expert and negative datasets.
We repeatedly apply the mutation operation to evolve the reward population R. each iteration, sample set of parent rewards R based on softmax distribution of their fitness, where exp(f(r))
the probability of selecting specific reward is given by EreR exp(s(r")) We apply the mutation operator to these parents to generate a new batch of candidate rewards. To maintain a high-qualig, population, we retain only the top N performing functions based on their fitness scores (from the combined set of current and newly mutated rewards). After K iterations, we return the single re- ward function with the highest fitness arg max (r)) This phase is detailed as function EVOSEARCH in Algorithm
Phase 3: Active Data Collection via Reinforcement Learning The optimal reward r* above obtained by inspecting existing demonstrations. further improve the reward, we can collect further demonstrations by training a policy лт,- using the current optimal reward r"; and use this policy to collect additional data D,,..
In more detail, employ PPO (Schulman al. 2017) to train policy in the environment of interest. As this process can be expensive, we use a predefined environment interaction budget added to the dataset of negative
N instead of training to convergence. The trajectories trajectories D- as they are likely to contain new edge cases the reward should consider. These are used further refine the reward population described in the preceding Sec. (Phase 2). This phase is presented as function DATAEXPANDRL in Algorithm []]
The final algorithm, presented in Algorithm I
consists of repeatedly performing Evolutionary Search over reward population R followed by data expansion using RL-trained policy.
Algorithm GRACE: Generating Rewards As CodE
Inputs:
function EYOSEEARCH(RR,S,.0. for k = 1 K do
D*: expert trajectories, D random trajec-
tories
Ruew for 1.. do
R: reward population
Parameters:
Sample r exp(f(r) r' <- m(r) // See Eq.
P: population size, N: RI. budget
M: data augmentation steps
Rnew - Roew U (r')
end for f Topp
procedure GRACE(D+, D
Keep best
S. { D+} S, {s D-}
rewards end for Return R end function
R = [L,M(S., S.. reward }
Reward Refinement.
for i 1.... M do R + EYOSEARCH(R, S.)
function DATAEXPANDRL(R)
D, + DAPAEXPANDRL(R)
7" <- arg max f(r) Train TT with PPO under budget N/M Collect new trajectories D,-
S, = D,, US,
end for
Return r" = arg max,eR f(r)
return D,,-
end procedure
end function
5

---

Published as a conference paper at ICLR 2026
Train
Test
1.0
1.0
Accuracy
0.8
Accuracy
0.8
0.6
0.6
0.4
0.4
20
60
80
20
40
60
80
Generations
Generations
AndroidControl
BabyAl
Figure 2: Fitness Number of generations. Evolution of train and test fitness across evolution generations, as defined by Algorithm]] [ for BabyAI and AndroidControl (multi-level settings). For BabyAI, provide 8 expert trajectories and negative trajectories for each task. For AndroidCon- trol, we provide 8-12 expert trajectories and ~880 negative trajectories total. Shading is standard deviation across 3 seeds. For these experiments, no online data is added beyond the initial trajecto- ries provided (M = 1).
Additional reward shaping When the reward function offline performance on D doesn't 't translate good online RL performance, we assume that the reward signal is poorly shaped, and additional refinement is required. In these cases, the LLM's info in Eq. 4 is augmented beyond low-fitness states to include full expert trajectories alongside the current reward value for each state. We then instruct the LLM to reshape the reward function, using expert trajectories reference, that it provides a signal that increases monotonically.
5 EXPERIMENTS
We empirically evaluate GRACE with respect its ability generate rewards that lead effective policy learning. Specifically, we aim to address the following questions:
Accuracy and Generalization: Can GRACE recover correct rewards, and with how many expert samples?
Policy Learning Performance: How does GRACE compare other IRL methods to online RL trained with ground-truth rewards?
Qualitative Properties: How well-shaped are the rewards produced by GRACE?
Interpretability and Multi- Task Efficacy: Does GRACE produce reusable reward APIs?
PPO
GRACE w/ GPT-4o
GRACE w/ Qwen3-Cuder-30B
GAIL w/
AIRL w/ 200 traj
200 traj
Hopper
2212 + 54
2143 + 80
2106 + 76
2056 + 92
2028 + 82
Walker
2675 + 292
2072 + 576
2229 + 600
1982 + 101
2108 + 293
Ant
6239 + 237
5707 + 210
6085 + 804
5521 + 674
4308 - 306
Humanoid
6455 + 302
5809 + 106
5921 + 301
6521 + 337
6512 + 291
Table 1: MuJoCo Results Average returns on 4 MuJoCo continuous control tasks. Average and standard deviation is reported across different seeds. The total number of required LLM calls to recover reward for each task averages at 2000 for both GPT-4o and Qwen3-Coder-300

---

Published as a conference paper at ICLR 2026
5.1 EXPERIMENTAL SETUP
To evaluate GRACE, we conduct experiments in three distinct domains: the procedurally generated maze environment BabyAl (Chevalier-Boise et al., 2018). which tests reasoning and generaliza- tion; the physics simulator MuJoCo (Todorov et al. 2012), which ensures GRACE works on con- tinuous, non-symbolic environments; and the Android-based UI simulator AndroidWorld (Rawles) al. 2024), to evaluate applicability real-world, complex problems. Across all experiments, don't Ton'tpon provide GRACE with any in-context information about the environment (such as a description or source code) ensure fair comparison against baselines.
BabyAI Our BabyA/ evaluation suite comprises 20 levels, including 3 custom levels designed to test zero-shot reasoning on tasks not present in public repositories, thereby mitigating concerns of data contamination. Expert demonstrations are generated using the BabyAI-Bot (Farama Foun- dation et al., 2025), which algorithmically solves BabyA/ levels optimally. We extend the bot to support our custom levels as well. For each level, we gather approximately 500 expert trajectories. Another 500 negative trajectories are collected by running randomly initialized agent in the envi- ronment. The training dataset consists of up to 16 trajectories, including both expert and negative examples. All remaining trajectories constitute the test set. For each dataset, we evolve the reward on the train trajectories and report both train and test fitness from Eq. (3). For this environment, we never augment the dataset with online trajectories (M = 1), we establish it's not necessary to recover the optimal reward.
The state is represented by (h, we array. The state is fully observable, with the first channel containing information about the object type (with each integer corresponding to a different object, such box, key, wall, or agent), the second channel contains information about the object's color and the third any extra information (e.g. agent direction, if is the door locked). The recovered reward function needs to operate on the state representation as (h, w. 3) array, but also in- clude a pixel image of the state in context for the LLM. For the classic GAIL baseline, we finetune llava- vwraonevii.o s on--qwen--0- 5b 5b-ov-hf on state images as the reward model. Given BabyAl is an environment that terminates success, we only consider the last state of expert trajectories as a positive state, we do this for both GRACE and the GAIL. baseline.
MuJoCo We conduct additional experiments on tasks from the classical MuJoCo continu- control suite (Todorov al. 2012): Hopper, Walker, Ant, Humanoid. These tasks demonstrate that GRACE also succeeds at reward design in continuous action and state spaces. We run all our MuJoCo experiments using the fully differentiable physics engine Brax (Freeman al. 2021) to speed up learning. Unlike the BabyA/ and Android experiments, in MuJoCo we augment (M 10) with new trajectories coming from the learner policy. The reward is only updated if the fitness is low on the newly added trajectories.
the dataset 10 times
Android To assess GRACE in high- sineh-immesional, real-world setting, we use the AndroidControl dataset (Rawles al. 2023 Li et al. 2024). which provides a rich collection of complex, multi-step human interactions across standard Android applications. The state space includes both raw screen pixels and the corresponding XML view hierarchy. Both representations are included in io-context in the mutation prompt, while the reward function only takes the XML representation as input.
From this dataset, we curate a subset of trajectories focused on the Clock application, where users successfully complete tasks such as "set timer 6 hours from now". These serve our positive examples. Negative samples are drawn from trajectories in other applications (e.g., Calculator, Calendar, Settings). Specificallyy we select 12 expert demos for the "set timer" task, and 8 each for the stopwatch tasks ("pause stopwatch" and "run stopwatch"). ). We use 80% of trajectories in the train set and the remaining for the test set. For this environment, online trajectories (M = 1), as we establish it's not necessary to recover the optimal reward. In this setting, recover multi-task reward and we give the option to the LLM to remove any states it judges as unnecessary to task completion.
never augment the dataset with
GRACE Parameters All parameters used across our experiments be found in Appendix A.4

---

Published as a conference paper at ICLR 2026
Task
PPO
GAIL
GRACE
Mean Episode Success
GoToReeBBllNaDiss
1.00
1.00
1.00
GoToRedBall
1.00
0.35
1.00
0.71
PickupDist
0.31
0.15
0.32
PickupLoc
0.21
0.00
0.26
0.70
GoTobjj
1.00
0.92
1.00
OpenDoorColor
1.00
0.98
1.00
0.69
OpenTwoDoee
1.00
0.37
1.00
Ground GRACE Reward
PlaceBetween (new)
0.09
0.01
0.09
0.68
OpenMatchingDooor (new)
0.79
0.20
0.35
0.0
0.5
1.0
1.5
2.0
Multi-task
0.95
0.31
0.92
Total Timesteps
1e5
Figure 3: Training Curves for Android- World Clock Tasks. Mean episode suc- cess over the 3 AndroidWorld clock tasks: ClockStopWatchPussee Verify, ClockStop- WatchRunning, and ClockTimerEntry.
Table 2: Success rates on selected BabyAL AI en- vironments. GRACE compared against PPO and GAIL.. GRACE uses expert trajectories per task, while GAIL uses 2000. Additional results are avail- able in Table
5.2 ANALYSIS
GRACE recovers rewards with high accuracy. We first examine whether GRACE evolutionary search (Phase 2) can successfully recover the underlying task reward from demonstrations alone. We evaluate this in two settings using BabyAl: (i) single-level setting, where the model infers task-specific reward, and (ii) a more challenging multi-level setting, where GRACE must learn a single, general reward function conditioned on both state and language goal.
In Figures 4 and we show that the reward accuracy consistently reaches 1.0 across all BabyAI tasks in both single- and multi-level settings, as well as on AndroidControl, A fitness of 1.0 corre- sponds to assigning higher values to all expert states than to negative states.
We further ablate two aspects of the algorithm. First, we analyze sample efficiency by varying the number of expert and negative demonstrations. Results on BabyAI (Figure/4a) show non-trivial per- formance even with single demonstration, with gradual improvement and perfect scores achieved using only eight expert trajectories. The number of negative trajectories also plays role, though to lesser degree: for example, accuracy of 0.95 achieved with just single negative trajectory, provided that sufficient expert trajectories are available (Figure[4))
Finally, we assess the robustness and efficiency of the evolutionary process. As shown Figure 2 in the multi-task setting GRACE reliably converges 100 generations (i.e., evolutionary search steps) and no additional trajectories from the learner agent (M 1). demonstrating the effectiveness of our LLM--rivee refinement procedure.
high-fitness reward function in fewer than
Comparison with Baselines: To validate the quality of the inferred reward function, we compare against two approaches. First, recovered by GRACE as well as the groundtruth sparse reward. Clearly, the latter should serve as an oracle, while does not benefit from dense rewards. As an IRL baseline, we compare against GAIL (Ho & Ermon 2016). GAIL is trained with large dataset of 2, 000 expert trajectories per task.
PPO (Schulman et al.) 2017) to optimise both the rewards
As shown in Tables and 2 GRACE consistently matches or outperforms baselines across all tasks with less training data. On several tasks. GRACE matches Oracle PPO with ground-truth rewards, whereas GAIL completely fails. This demonstrates that the interpretable, code-based rewards from GRACE are practically effective, enabling successful downstream policy learning. To ensure fair comparison, baseline and GRACE agents are trained using the same underlying PPO implementa- tion, agent architecture and hyperparameters. Performance measured by the final task success rate after le7 environment steps. Crucially, no extra information or environment code is provided in context to GRACE. Similarly, we evolve separate reward function for each task in the AndroidCon- trol dataset matching tasks present in the Clock AndroidWorld tasks: ClockStop WathhWWWWhhhsseeeee Verify, and ClockTimerEin The training curves for all tasks (averaged) are reported in Figure
8

---

Published as a conference paper at ICLR 2026
8
0.59
0.70
0.81
1.00
Negative Trajectories
00
Test Accurocy
0.8
0.46
0.69
0.88
0.99
Test Accura
2
0.46
0.69
0.89
0.97
0.4
1
0.61
0.73
0.86
0.95
0.0
Expert Trajectories
Expert Trajectories
(b)
Figure 4: Fitness vs Number of Expert Trajectories. The accuracy is computed on test dataset after obtaining maximum fitness on training data with corresponding number of expert and negative training trajectories. (a) Performance on all 20 BabyAl tasks. (b) Aggregate accuracy across 20
BabyAl tasks.
1.0
agent_pos
aan
Call Number
15-
0.9
Call Number
20
go. reward
target
Accuracy
15
10
10
5
Reused Modules New Modules
0.6
10
20
40
50
60
30
40
50
Generations
Generations
Figure 5: Module and function reuse across generations On the left, show at each generation step the number of newly created modules and number of existing and thus reused modules from prior rewards, contrasted with the accuracy in the reward population. On the right, show number of times module are being re-used, for select set of modules.
MuJoCo Performance In our MuJoCo evaluation, GRACE leverages an LLM- LLM-generated debug function analyze distributional differences between expert and negative trajectories, computing key statistics such as mean, standard deviation, percentiles, and Cohen's 's d for each state feature. While this statistical analysis proves effective for Nower-dimensione. tasks, the method still struggles with the Humanoid environment. We hypothesize that this performance drop stems from the LLM's struggle effectively parse and reason the extensive numerical arrays required represent high-dimensional states in text-based format.
GRACE generates well shaped rewards: We demonstrate GRACE's ability produce well- shaped rewards that accelerate learning. For challenging, long-horizon tasks like OpenTwoDoors. correct but unshaped reward lead to local optima where the agent gets stuck (Figure "Gen 1"). By explicitly tasking the LLM to introduce shaping terms, GRACE refines the reward to provide denser learning signal. As shown in Figure 6 this targeted shaping dramatically improves both the final performance and the speed of learning, allowing the agent to solve the task efficiently. This confirms that GRACE not only finds what the goal but also learns how guide an agent towards it.
GRACE Code Reuse: A key advantage representing rewards as code the natural emergence of reusable functions that collectively form domain-specific reward library. We study this phe- nomenon in the multi-task BabyAl setting (Figure In the early generations of evolutionary search, GRACE actively generates many modules to explore alternative reward structures. After gener- ation 10, the rate of new module creation drops sharply. At this point, GRACE shifts toward reusing the most effective, high-level modules it has already discovered.

---

Published as a conference paper at ICLR 2026
1.0
Reward on Expert Trajectory
Mean Episode Success
1.50
Gen 1
Gen 2
Gen 3
0.6
Reward
1.00
Gen
0.75
Gen
0.2
Gen 3
0.25
Gen 4
0.00
0.0
0.2
0.4
0.6
0.8
1.0
0
8
10
12
Total Timesteps
1e7
Timestep
Figure 6: Shaping Using the default reward recovered by GRACE occasionally leads to failure in learning the correct behavior due poor shaping. Through the targeted shaping in Phase 3, significantly improve final performance and speed of learning.
To further illustrate this reuse, Figure (right) shows call counts for selected set modules within the evolving reward API. For instance, the Goal module, which summarizes of goals, is initially used sparingly but becomes heavily invoked following a code refactor at generation 30. Likewise, the agent_pos function is reused least five times after its introduction. These trends demonstrate that GRACE progressively builds a reward library that supports efficient multi-task generalization.
DISCUSSION
Limitations A key limitation of GRACE is its limited scalability to high-dimensional state spaces for evolving reward functions. While code offers interpretability, neural networks are inherently better suited for processing high-dimensional observations (such pixels), as they excel at learning distributed representations directly from raw sensory data, a task where symbolic feature extraction often proves brittle or infeasible.
Conclusion We introduce GRACE, a novel framework that leverages LLMs within an evolution- ary search address the critical challenge of interpretability in IRL. Our empirical results demon- strate that by representing reward functions executable code, we can move beyond the black-box models of traditional IRL and produce rewards that are transparent, verifiable, and effective in RL learning. We show that GRACE successfully recovers accurate and generalisable rewards from few expert trajectories, in stark contrast to deep IRI. methods like AIRL. This sample efficiency sug- gests that the strong priors and reasoning capabilities of LLMs provide powerful inductive bias. Furthermore, we demonstrate the framework's 's practical utility by applying to the AndroidWorld environment, showing that GRACE can learn rewards for variety of tasks with real-world applica- tions directly from user interaction data.
REFERENCES
Pieter Abbeel and Andrew Y. Ng. Apprenticeship learning via inverse reinforcement learning. In Proceedings of the Twenty-First International Conference on Machine Learning (ICML), pp. 1-8.
ACM, 2004.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoe V. Le, and Charles Sutton. Program synthesis with large language models. arXiv preprint arXiv:2108 07732, 2021.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Michael Pavlov, Alethea Power, Lukasz Kaiser, Miljan Bavarian, Clemens Winter, Phil Tillet, Felipe Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert- Herbert-Voss, William Guss, Alex Nichol, Igor Paino, Nikolas Tezak,
10

---

Published as a conference paper at ICLR 2026
Jie Tang, Igor Babuschkin, Suchir Balaji, Suyog Jain, William Saunders, Christopher Hesse, Mark Carr, Aitor Lewkowycz, David Dohan, Howard Mao, Lily Thompson, Erica Frank, Joshua Chen, Victor Butoi. David Hernandez, Liane DasSarma, Maxwell Chan, Mateusz Litwin, Scott Gray, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Evaluating large language models trained code. arXiv preprint arXiv:2107.03374. 2021.
Maxime Chevalier-Boisvert, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning. arXiv preprint arXiv:1810.08272 2018.
Paul F. Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In Advances in Neural Information Processing Systems (NeurIPS), pp. 4299-4307, 2017.
Stephen A. Cook. The complexity of eecorem-proving procedures. In Proceedings of the Third Annual ACM Symposium on Theory of Computing, STOC '71, pp. 151āÄŞ158 New York. NY, USA, 1971. Association for Computing Machinery. ISBN 9781450374644, doi: 10.1145/800157. 805047. /tRt:s/t::///at http:////0....o11..1111111101110 805047
Nicola Dainese, Matteo Merler, Minttu Alakuijala, and Pekka Marttinen. Generating code world models with large language models guided by monte carlo tree search. Advances in Neural Infor- mation Processing Systems, 37:60429-6044 2024.
Agoston E. Biben and J. E. Smith. Introduction to Evolutionary Computing. SpringerVerlag. 2003a. ISBN 3540401849.
Agoston E. Eiben and James E. Smith. Introduction to Evolutionary Computing. Springer, 2003b.
Maxence Faldor, Jenny Zhang, Antoine Cully, and Jeff Clune. Omni-epie: Open-endedness via models of human notions of interestingness with environments programmed in code. In Interna- tional Conference on Learning Representations, 2025. URL https::/
Farama Foundation, Maxime Chevalier- Bolun Dai, Mark Towers, Rodrigo Vicent,, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid: Modular customizable reinforcement learning environments. hhtpps: /tap:::///itiie htg.gtttubb...m com/
Frrama-Founndtttoon//iiiii
2025. Accessed: 2025-09-24.
C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax differentiable physics engine for large scale rigid body simulation, 2021. URL. http:
//gi t
Daniel Fried, Joshua Ainslie, David Grangier, Tal Linzen, and Dani Yogatama. Incoder: A genera- tive model for code infilling and synthesis. In International Conference on Learning Representa- tions (ICLR), 2022.
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse rein- forcement learning, 2018. URL htttpss//ar....oor///////111 11248
Kurt Godel. Letter john von neumann, 1956. URL wtt./t/tps//eeoooooo.roeeeee http:///..ccooeettttteeeee//aaa/ content
David E. Goldberg, Genetic Algorithms in Search, Optimization and Machine Learning. Addison-
Wesley Longman Publishing Co., Inc., USA, 1st edition, 1989a. ISBN 0201157675.
David E. Goldberg. Genetic Algorithms in Search, Optimization, and Machine Learning. Addison- Wesley, 1989b.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in Neural
Information Processing Systems, volume 29, 2016.
11

---

Published as a conference paper at ICLR 2026
Martin Klissarov, Pierluca D'Oro, Shagun Sodhani, Roberta Raileanu, Pierre-Luc Bacon, Pascal Vincent, Amy Zhang, and Mikael Henaff. Motif: Intrinsic motivation from artificial intelligence feedback. arXiv preprint arXiv:23 2023.
Martin Klissarov, Devon Hjelm, Alexander Toshev, and Bogdan Mazoure. On the modeling capabil- ities of large language models for sequential decision making. arXiv preprint arXiv:2410.05666, 2024.
Raymond Li, Loubna Ben Allal, Yacine Jernite Zi, Denis Kocetkov, Chenxi Mou, Aleksandra Piktus, Laura Weber, Wenhao Xiao, Jihad Bibi, Stella Biderman, et al. Starcoder: may the source with you! arXiv preprint S00.:::30000066... 2023.
Wei Li, William Bishop, Alice Li, Chris Rawles, Folawiyo Camphell-Ajala, Divya Tyamagundlu, and Oriana Riva. On the effects of data scale on computer control agents. arXiv e-prints, pp- arXiv-2406, 2024.
Yecheng Jason Ma, William Liang, Guanzhi Wang, De-An Huang, Osbert Bastani, Dinesh Jayara- man, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Fureka: Human-level reward design via coding large language models. arXiv preprint arXiv: Arxiv-2310. 12931, 2023.
Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, et al. Self-refine: Iterative refinement with self-feedback. Advances Neural Information Processing Systems, 36:46534 46594, 2023.
Christoph Molnar. Interpretable machine learning. Lulu. com, 2020.
Andrew Y. Ng and Stuart J. Russell. Algorithms for inverse reinforcement learning. In Proceedings of the Seventeenth International Conference on Machine Learning (ICML), pp. 663-670. Morgan Kaufmann, 2000.
Erik Nijkamp. Richard Pang, Hiroaki Hayashi, Tian He, Baptiste Roziere, Canwen Xu, Susan Li, Dan Jurafsky, Luke Zettlemoyer, Veselin Stoyanov, and Hyung Won Chung. Codegen: An open large language model for code with multi-turn program synthesis. In International Conference on Learning Representations (ICLR), 2022.
Alexander Novikov, Ngån Vü, Marvin Eisenberger, Emilien Dupont, Po-Sen Huang, Adam 7solt Wagner, Sergey Shirobokov, Borislav Kozlovskii, Francisco JR Ruiz, Abbas Mehrabian, et al. Alphaevolve: A coding agent for scientific and algorithmic discovery. arXiv preprint arXiv:2506.13131, 2025a.
Alexander Novikov, NgAén vÅr, 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: A coding agent for scientific and algorithmic discovery, 2025b.
/506. 13131
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kel- ton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. arXiv preprint arXiv:2203.022 02055, 2022.
Emilio Parisotto, Jimmy Lei Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and trans- fer reinforcement learning. In International Conference on Learning Representations (ICLR),
2016.
Martin I. Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Nathan Ratliff, J. Andrew Bagnell, and Martin Zinkevich. Maximum margin planning. In Proceed- ings of the 23rd International Conference on Machine Learning (ICML), pp. 729-736. ACM, 2006.
12

---

Published as a conference paper at ICLR 2026
Christopher Rawles, Alice Li, Daniel Rodriguez, Oriana Riva, and Timothy Lillicrap. Android in the wild: A large-scale dataset for android device control, 2023. URL hRtp:ttts://xxii... / abs abs/2307 10088
Christopher Rawles, Sarah Clinckemaillie, Yifan Chang, Jonathan Waltz, Gabrielle Lau, Mary- beth Fair, Alice Li, William Bishop, Wei Li, Folawiyo Campbell-Ajail et al. Androidworld: A dynamic benchmarking environment for autonomous agents. arXiv preprint arXiv:2405.14573 2024.
Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Matej Balog, M. Pawan Kumar, Emilien Dupont, Francisco J. R. Ruiz, Jordan S. Ellenberg, Pengming Wang, Omar Fawzi, Pushmeet Kohli, and Alhussein Fawzi. Mathematical discoveries from pro- gram search with large language models. Nature, 625(7995):468 475, 2024a. doi: 10.1038/
541586-023-0692446
Bernardino Comera-Paredess Mohammadamin Barekatain, Alexander Novikov, Matej Balog, M Pawan Kumar, Emilien Dupont, Francisco JR Ruiz, Jordan S Ellenberg, Pengming Wang, Omar Fawzi, al. Mathematical discoveries from program search with large language models. Nature, 625(7995):468 475, 2024b.
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies scalable alternative to reinforcement learning. In Proceedings of the 36th International Confer- ence on Machine Learning (ICML), 2017a.
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ily Sutskever. Evolution strategies as scalable alternative to reinforcement learning. 2017b.
1703. 03864
Silvia Sapora, Gokul Swamy, Chris Lu, Yee Whye Teh, and Jakob Nicolaus Foerster. Evil: Evolution strategies for generalisable imitation learning, 2024.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy 2017.
optimization algorithms. arXiv preprint
Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in Neural Information Processing Systems (NeurIPS), volume 12, pp. 1057-1063, 1999.
Gokul Swamy, Sanjiban Choudhury, J Andrew Bagnell, and Steven Wu. Of moments and match- ing: A game-theoretic framework for closing the imitation gap. In International Conference on Machine Learming, pp. 10022-10032. PMLR, 2021.
Umar Syed and Robert E. Schapire. A game-theoretic approach apprenticeship learning. Advances in neural information processing systems, 20, 2007.
Yee Whye Teh, Victor Bapst, Wojeiech M. Czarnecki, John Quan, James Kirkpatrick, Raia Had- sell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), pp. 4509-4509, 2017.
2017.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-50000000 2033. IEEE, 2012. doi: 10.1 10111099RRO000 2.6386109,
David Venuto, Mohammad Sami Nur Islam, Martin Klissarov, Doina Precup, Sherry Yang. and Ankit Anand. Code reward: Empowering reinforcement learning with VLMs. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.). Proceedings the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 49368 49387. PMIR, 21-27 Jul wtwpssmmr.prrss./emmmmmm
2024a. URI.tttp... ptttt:::/proceedl ngs ml pr..pres/vvee v2.5ss/v223///./ http//////////e22////
html
David Venuto, Sami Nur Islam, Martin Klissarov, Doina Precup, Sherry Yang, and Ankit Anand. Code as reward: Empowering reinforcement learning with vlms. arXiv preprint arXiv:2402.04474 2024b.
13

---

Published as a conference paper at ICLR 2026
Yuxiang Wang, Yuchen Lin, Dongfu Jiang, Bill Y. Chen, Xiang Shen, Jidong Zhao, Xiang Yu, Chen Li, Xiao Qin, and Jie Sun. Lim-blender: Ensembling large language models with pairwise ranking and generation. In Proceedings of the 6lst Anmual Meeting of the Association for Computational Linguistics (Volume I: Long Papers). Association for Computational Linguistics, 2023.
Aaron Wilson, Alan Fern, and Prasad Tadepalli. Multi-task reinforcement learning: A hierarchical bayesian approach. In Proceedings of the 24th International Conference on Machine Learning (ICML). pp. 1015-1022. ACM, 2007.
Wenhao Yu, Nimrod Gileadi, Chuyuan Fu, Sean Kirmani, Kuang-Huee Lee, Montse Gonzalez Are- nas, Hao-Tien Lewis Chiang. Tom Erez, Leonard Hasenclever, Jan Humplik, al. Language to rewards for robotic skill synthesis. arXiv preprint arXiv:2306.0888 2023.
Abhay Zala, Jaemin Cho, Han Lin, Jachong Yoon, and Mohit Bansal. Envgen: Generating and adapting environments via Ilms for training embodied agents. In Conference on Language Mod- eling (CoLM), 2024.
William Zankner, Rohan Mehta, Eric Wallace, Jack Fitzsimons, Yang, Alex Mei, Daniel Levy, William S. Moses, and Joseph E. Gonzalez. Critique-out-loud reward models. 2024. URL
wttp://.rraaab////////// 11791
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging llm-as-a-judge with mt-bench and chatbot arena. Advances in neural information processing systems,
2023.
Andy Zhou, Kai Yan, Michal Ciaapentokh-Rothmaee Haohan Wang, and Yu-Xiong Wang, Lan- guage agent tree search unifies reasoning acting and planning language models. arXiv preprint arXiv:2310.o0 044666 2023.
Brian D. Ziehart, Andrew I.. Maas, J. Andrew Bagnell, and Anind K. Dey. Maximum entropy inverse reinforcement learning. In Proceedings of the Twenty-Thirg AAAI Conference on Artificial Intelligence (AAAI), pp. 1433-1438. AAAI Press, 2008.
14

---

Published as a conference paper at ICLR 2026
A APPENDIX
A.1 ADDITIONAL ONLINE RESULTS
Task
PPO
GAIL
GRACE
w/ 2000 trajs
w/ 8 trajs
GoToRedBalNNoDist
1.00
1.00
1.00
GoToRedBall
1.00
0.35
1.00
PickupDist
0.31
0.15
0.32
PockupLoc
0.21
0.00
0.26
GoToObj
1,00
0.92
1.00
OpenDoorColor
1.00
0.98
1.00
OpenTwoDoors
1.00
0.37
1.00
OponRedDoor
1.00
1.00
1.00
GoToObjS4
1.00
1.00
1.00
Go Blll
0.96
0.40
0.99
Go GoTedBBallGrry
0.97
0.77
0.99
Pickup
0.10
0.00
0.09
Open
0.30
0.18
0,22
OpenRedBlue
1.00
0.96
0.98
OpenDoorLoc
0.39
0.40
1.00
TooooScc888TT
0.64
0.39
0.97
GoToDoor
0.74
0.37
0.99
SortColors (new)
0.00
0.00
0,00
PlaceBetween (new)
0.09
0.01
0.09
OpehhatchingDooor (new)
0.79
0.20
0.35
Multi-task
0.95
0.31
0.92
Table 3: Success rates additional BabyAl environments. The performance of method, GRACE, is compared against two key basclines; PPO, trained on the ground-truth reward, and GAIL, trained using 2000 expert trajectories per task. GRACE's performance is evaluated with expert trajectories per task to demonstrate its high sample efficiency. All values represent the final success rate at the end of training. We don't report GAIL's performance 8 expert trajectories as it is near 0 for most tasks.
A.2 EXTENDED DISCUSSION AND FUTURE WORK
GRACE's reliance on programmatic reward functions introduces several limitations, particularly when compared to traditional deep neural network based approaches. These limitations also point toward promising directions for future research.
Input modality While generating rewards as code offers interpretability and sample efficiency, struggles in domains where the reward depends complex, high-dimensional perceptual inputs. Code is inherently symbolic and structured, making it less suited for interpreting raw sensory data like images audio. For instance, creating programmatic reward for task like "navigate the cat" is non-trivial, as "cat" is difficult visual concept. NNs, in contrast, excel at learning features directly from this kind of data. Programmatic rewards can also be brittle: small difference in the environment violating a hard-coded assumption could cause the reward logic to fail completely, whereas NNs often degrade more gracefully.
Data Quantity GRACE demonstrates remarkable performance with very few demonstrations. This is strength in data-scarce scenarios. However, it is limitation when vast amounts of data are available. Deep IRL methods like GAIL are designed to scale with data and may uncover complex patterns from millions of demonstrations that would be difficult to capture in explicit program. While GRACE's s evolutionary search benefits from tight feedback on a small dataset, it is not clear how effectively it could leverage massive, complex dataset.
15

---

Published as a conference paper at ICLR 2026
Failure Cases Although GRACE is highly sample-efficient, it can still fail. For example, in the BabyAl-Open task. GRACE often proposed a reward that didn't t take into account the order in which the doors were being opened. Similarly, in the new BabyAl-SoriColors task, it would sometimes return a reward that only accounted for picking up and dropping both objects, without paying attention where they were being dropped. While these errors be easily fixed by providing a relevant negative trajectory or by treating all learner-generater states as negative trajectories, they highlight that GRACE can still misinterpret agent's intent based on expert demonstrations alone.
Hybrid Approaches These limitations can substantially mitigated by extending the GRACE framework to incorporate tool use, combining the strengths of both systems. The LLM could be granted access to a library pre-trained models (e.g., object detectors, image classifiers, or seg- mentation models). The LLM's task would then shift from writing low-level image processing code to writing high-level logic that calls these tools and reasons over their outputs. A final direction involves generating hybrid reward functions that part code and part neural network. The LLM could define the overall structure, logic, and shaping bonuses in code, but instantiate a small, learn- able NN module for specific, difficult-to-program component of the reward. This module could then be fine-tuned using the available demonstrations, creating a reward function that is both largely interpretable and capable of handling perceptual nuance. By exploring these hybrid approaches, future iterations of GRACE could retain the benefits of interpretability and sample efficiency while overcoming the inherent limitations of purely programmatic solutions in complex, perception-rich environments.
16

---

Published as a conference paper at ICLR 2026
A.3 NEW BABYAI LEVELS
To evaluate the generalization and reasoning capabilities of GRACE and mitigate concerns of data contamination from pre-existing benchmarks, designed three novel BabyAl Al levels.
PlaceBetween The agent is placed in a single room with three distinct objects (e.g., red ball, green ball, and a blue ball). The instruction requires the agent pick up specific target object and place it on empty cell that is strictly between the other two anchor objects. Success requires being on the same row or column as the two anchors, creating a straight line. This task moves beyond simple navigation, demanding that the agent understand the spatial relationship "between" and act upon a configuration of three separate entities.
OpehhnaccingDDor This level is designed to test indirect object identification and chained infer- ence. The environment consists of a single room containing one key and multiple doors of different colors. The instruction is "open the door matching the key". The agent cannot solve the task by simply parsing an object and color from the instruction. Instead, it must first locate the key, visually identify its color, and then find and open the door of the corresponding color. This task assesses the agent's s ability to perform simple chain of reasoning: find object A, infer property from it, and then use that property identify and interact with target object B.
The environment consists of two rooms connected by door, with red ball in one room and a blue ball in the other. The instruction is a compound goal: "put the red ball in the right room and put the blue ball in the left room". To make the task non-trivial, the objects' initial positions are swapped relative to their goal locations. The agent must therefore execute a sequence of sub-tasks for each object: pick up the object, navigate the other room, and drop it. This level tests the ability to decompose complex language command and carry out a plan to satisfy multiple, distinct objectives.
SortColors
17

---

Published as a conference paper at ICLR 2026
A.4 HYPERPARAMETERS
Table 4: Hyperparameters for Training BabyAl with PPO
Parameter
Value
Base Model
Gamma 0,999
Learning Rate
3e-5
Entropy Coef
le-5
Num Envs
10
Num Steps
64
Episode Length
100
PPO Epochs
Num Minibatch
Table 5: Hyperparameters for Training AndroidWorld
Parameter
Value
Base Model
-0ee2.5-VL33B-Insttee
LoRA Rank
512
LoRA Alpha
32
LoRA Dropout
0.1
Critic Hidden Size
2048
Critic Depth
4
Gamma
0.999
Learning Rate
3e-5
Entropy Coef
0.0
Num Envs
16
Num Steps
Episode Length
20
PPO Epochs
Num Minibatch
2
Table 6: Hyperparameters for GRACE Evolution
Parameter
Value
Population Size
20
Elite
Num Generations
100
Include expert trajectory chance
0.25
Incorrect state only chance
0.5
Expert state only chance
0.75
Model
gpt-4o
18

---

Published as a conference paper at ICLR 2026
A.5 EVOLUTION EXAMPLES
def www.aa.ee............ _pw.seaaaelooooeeeeeeee wow.oooooooooteeeeeeeeoooo Lext (Lex:: (pLional(st:]) if text is None:
(ptional(in:):
return None
colour_ words: Diet [ste, InL] - "red": D,
"green":
"blue"s 2,
"yollow" 6w. +r
parple" 3,
"ye"lows low":
*orange": 5,
old mapping
12 "gzey": 5, alias for the observed colour zode in the trajectery
"gray*: 5,
15
Lowez Lext, Lext,
16
for word, code in colous_words, icena
if word centrrowes lowe2:
return code
19
return None
21
_gacse_goel__yee (Lext: Optional (st:]) stz: if text is None:
23
return "key"
25
ixt (ower!
26
if "ball" in Ext:
return "bal
29
return "box"
30
return "key"
Figure 7: GRACE iteratively refines the initial BabyAl reward function (iteration 0) to handle unseen entities (iteration 10). Using execution traces, the agent fixes its color code mistake and adds a new box entity.
Irom future Import sanotations
import ze
typing import Optional, Tuple
import numpy 85 np
COLOR2ID :
*red":
10
*green": 1,
*blue":
12
*purple": 3,
"yellow":
14
"grey": 5,
15
"gray": 5,
Us spelling
17
QBJECT21D
19
*smpty": 0,
wall** 1,
21
#floor:: 2,
22
"door:: 3,
*kay":
24
*all"" 8,
"axx":
36
*agent": 10,
27
28
BTR2VEC: dict [int, Tuplelint, int]]
30
D: (1, 0), south
31
1: (0, 21, cast
(1, C)
33
3: (0, -1), t west
34
19

---

Published as a conference paper at ICLR 2026
_earse_gee (extra
Tuple(int, Uptional
37
"""Betuzn (objjctt_id) colour_id) pazaed Irom textre_intos,"
38
if not extra)_infoo inno: (xoProrr("xxtr__ must speciiy the target, e.g- the red ball
40
tokens fofffinnlllii [e-" Matral_info.lee
42
obj_id: Optionel Nene
col _ol_ optional(1ng None
for tok in tokens: +f
46
ir tok in COLORZID and col_id la None:
47
col_id COLOR2ID
in OBJECT2ID ebj_id
ob1_'d [FTFC
50
eet None and bok in S9LOR215T
obj_id None:
53
raine Valuegrror(
#"Could not parse target object extra_info- (extra_infol"..
35
56
return obi_id, col_id
colour may be None (xild-card)
58
Reward:
Sncoues when defini:ion ion (single-step) dense reward) :
103.0 aăş rsf cell front the agent soithere adjacent (according
clesest target object (Honhottan destance eR + direction szored in the third observatien. channel) conteins has disappesred from the observable grid lpicked upt.
64
65
matching
target
66
Kanhetten distance to the
68
bergety sippedd-at object.
69
<1.0 2Aş shaping revord 1/ (d+1) otherwise.
70
0.0 añs if either the agent or la matching) target is out-of-vees viewe not visible.
71
72
The implementation is modular new gos"
73
types can handled by extending the 0BUECT/COLOR lookup tables oz by repiacing +success predicater.
76
SUCCESS 100.0
78
def init (self, extral Linfo (oti str None)
79
self-tgt_cbhjic seif self target_org salg._ lexteaa
sp.ndarray) foott I anable direct-call
solf. www.frrworrdddoo (sttate)
84
def seward_fn(sele state: up.ndarray) >> float:
85
"*"state: (E, 3) n*n agent_pos self. aind_agt (state)
36
87
if agent is None:
0.0
89
90
4 of 341 target aojee=s-atii vigible tgt_maak = (statel:, 3, 0) statelt, ;, 1] solf.tgr_col_it
91
6
94
# pot Hokk
objjct Posject-gone--o wow.-o)-----oo paapaet--ggoo---pooo-------- casried
self sUCCEsS, RENARD
97
pooitiinonn:
20

---

Published as a conference paper at ICLR 2026
target_posicions self.__ felli__fin__taageet targeets (state)
if target_.ossiions.skt 0:
102
No matching Lazget in view -> no shaping.
return 0.0
104
105
Success predicale 3Aş Lazget must be directly front of the agenz. ir wt______trge__nn_eeeoeo tagent pos, poa,
106
107
state):
return
109
Shaping: inverse distance
target.
111
op.aba Cetarret_ppsstiinnt agent_ post (oxiss===) min_dist int (dists.min())
113
114
if min_diat <=
self, SUCCESS RENDRR
return 1.0 (1.0 min_dist)
117
estaticmethod
118
def _find_aaaee agant (state: up.ndarrey) -> Optional (np.ndarray):
119
axmpetten
cest os "*"Locate the fieet agent
pixel found,
(row, col)
..*
abssent. *..
120
locs np (stggeee ", 0] OBJECT2ID (aggent11 if cecs.size 0:
122
reture None
134
def _rwind_aargesseeeeeee_ee scate: np-ndarrayl -> np.ndarray:
125
"*"Return
21 array of rew/col pos' tona matching targeta.
126
obj_mask - statel;, :, 0] -- self, telf__taarge_objj...
127
if selg__iarget_cooour__ is not None:
128
col_maak statel;, ;, 1J self, Larggatttootoooo___ttt
129
mask obi_maak col_mask
130
e1se:
mask sbj_maak
132
return ap.orguhere (mask)
133
def it__arget__ta in__nn_ff front (self, agant_poa: np.ndarray, np.ndarrayl bool: 134
135
***Rcturn #ruue ifs the cell directly in front of the agent matches target.
136
rox, col agent_pos
agent_dir
staatesow, col, 21)
138
drow, d:el DIR2VBC. ge: gegeett
0)) 4 defsult to scuth if unknown c_w col zew drow, col deol
139
140
bounds cannot success.
142
if not (0 f_rew shape and <= f_col state. shape [1]):
143
return ralse
144
145
Check object id if state:f_row, f_ool, 0] self, Largge____gge__b______ iid:
146
148
Check colour colour spedified.
150
if
151
self._tareee ____aaret__ooou________ ld is no: None
152
[tato( f_ool, self. wallf.tttrget__oc our_id
153
):
154
return False
return
Figure 8: Example of code evolution across many generations.
21

---

Published as a conference paper at ICLR 2026
A.6 GENERATED REWARDS
A.6.1 ANDROID: SET TIMER REWARD
Import Json
Import re
Izom typing import Optional, Tuple
GENPRIC NORMALISATTON KRLPERS
def [ab__eelectedsstacee str, label: str) bool:
10
patiern testeent_descrsspiiolleess "aser)s"" (Maae)"""""
=E
12
return bool (re.search (pattern, state, re.T))
14
stime__ab__eleeceeiiee stx)
16
return _lab_selecced(etate, "Time:")
17
18
def _ourton_visillettt (state: label: str) -> bool:
19
return bool
se.search drontent_descriptiie
""errke"" atate, re.T)
21
23
def
(state: stz) bool:
24
return bottor (state, "Pause")
26
def _t _timer_keppad mode (state: scr) +> hool:
bool
"bbldii, 20t.ee 2. stace))
28
30
TIMER DURATTON PARSTNG
31
33
def _aaree______ rearss_requested__o time (text: str) -> Int:
*"*Parses natural language (e.g. minutes') seconds.
35
hours replacee"c""" ")
text text.
seconds
37
38
Check explicit hours, minutes, seconda for mult
40
(sdr) \ashour", 3600),
(Ad+) teminntee,,,"" 60),
42
(r" (1d+) \sxsecond", 1),
11
44
re.search (patt, text, re.I)
45
if m:
(oogrrup
47
if mult -- 3600:
3600
49
clif mult == 60:
minutes
51
else: seconds val
52
54
Fallback: 15 "min" or zaw aumbez found
minutes m re. search (i* (1d+) Noommnnm text, re.T! 3:
0:
56
57
58
minutes int (1))
59
else:
seacch (di)), text)
61
if m:
(m.group
63
64
total hours 3500 minutes 60 seconds
65
return total if total else 60
66
68
STATE EXTRACTION
70
71
def _extrect_ttiiee wenettttttme_eoeee Lomee__comooenttseeee (state:: str) Optional [Tuple int, int]]:
12
""*gxtracts (AB, 58) from the Android u1 hierarchy state, *na
73
Text format minutes seconds" re. (e" (1d+) Namminutor (-- (Ad+) \aesaconda", state, re. IGNCRECASE)
75
m:
22

---

Published as a conference paper at ICLR 2026
77
minutes (m. group (11)
78
seconda = int (m. gn.ggo (21) return (0, minutes, seconda)
79
80
81
Case 2: Compect format "Ih 05m 20 re.ceerch(s" (1d)) m:
state, (ONCRECCC)
83
(m.group (1))
85
minutes = Int [m. groppp
86
seconds int (m.groop (31)
87
return (hours, minutes, seconds)
88
Case "MX:SS" format (excluding standard time 1ike 12:30 PM
90
for mrnatacc in ze, finditer (s" (1) (Ad(21) (P1N8+ [AaPp] [Mr])", " state):
mn, 56 (mm_match. (tm_matcch.gropp(1) int (mm_natch.groopp22 it not (0 <= 88 63):
92
93
continue
context state (mntmmschhccc end() mmmmetch.ead() (o]eee "minute" in centext or "timar" in context or "ramsining" in context:
95
mn,
97
Digging through if not _twier_reea___eeeee (eteee___dd_eelleeeeeedddd.. return None
hiererchy Timer active
99
100
102
tokers = re. findall jr' "toxt" (t.strip() for tokens]
state)
104
Look RH:MS:SS pattern in For i (Loken)) 4):
107
takens;i]]
109
and tokena(1 + 1]
1e.tuullmatcca:aaa::31 and tokens [1 + 31 -- *;"
Lokena 2))
111
re. f
tokens(i 4))
114
int (tokens (il)
m_val Itokens
116
a = int (tokens [1 4]1 if return th, m_val,
118
Look for WX:S8 pattern in tokens for i in sange (len (tckens) 2):
121
123
re. re. (t Lokens
*:*
125
and re. fullmathh 1= \d/21", tokens(i - 2])
mLval int (tokens(i))
128
s_val int Itokens - 2)) if e_val 60: zeturn 10, m_val, s_val)
130
132
return None
daf _sa_a_jjoon_____________________ sarfe_jjoo_dduppp (abt)
135
try:
fson.dumpa tob), essure_eecii=Falee)
137
except Exception:
return json..umngg (1 ("erroo": sdbbaggserrisliaattio failed"1)
139
141
RENARD CLASS
142
144
class Reward:
rexard function specifically setting Google Timors."*a
146
def Limit (self, extra_info: Optional:sti [str] None): seif.rew_instr:: str extra_info e* seil..nstrruttion st: selt.caw_inett..ovee
149
151
self.task.type "set_timer"
153
Goal parsing
selffgooi_secontee parse saasse_reueesedd_cc time (self, inst.iittttti
seae.goal_seconds 3603
156
self.goal_ins secosdd // 3600, zem // 60, rem. & 60)
23

---

Published as a conference paper at ICLR 2026
self.goai_
159
self._t =
161
def seward_fn(sell, state: stz). > flost:
self._t
: return 100.0
164
166
Since are only doing sec timer zeturn rewar_ (state)
def @roug_fn(sel: state: ste) +> strr
171
"step": self._t,
"task_typa": sallf.task_tpp
173
"goal "goal_achieved"; sel. gel2.-gaa "gosi_hms": sel5.goal_hms,
176
return _aaaa__ json___oon_duupp (dbg)
178
def _reeard Limer (aelf, state: sL:) float:
181
# Check i5 we are on the cerrect tab if _timer_tab_see tab_selsctae(sttttt) reward 0.2
183
185
Extract current digits entered into the timer current_val _extract__ingg comerpociimtssomoenntttt (state) if eurzent_val is None:
187
188
return min (reward,,, 0.99)
190
eur_hh, our_me, cur_ss ourrent_va
192
Normalize to digit strings for comperison (strip leading zecos) srreet_ddiit_ssriiiee f"/cur_hh:02d cur_mm::0d) (asr__ss:02)]) "lattii ("0*)
-istsip
if cuszent_digit_stiiie =- **;
195
current_digee current_digit_atrring *g"
197
goa
I"(aelf.goal_imm hma [0] :02d] [ael..gosl_hmg [1] [02d] [sel gomm [2] (2]]000 :02
Latrip ("0")
198
if geel sel__digt__stittee == "*;
goal_cigit_atrrng
201
Exact natch check
current_die aigit_string
goal_dig' t_string:
203
We Lreat exact digit entry as success teven if no: started yet) values.
timer already running reture 100.0
205
Partial match revard (dense reward for typing correct digits) aetching_digits in range min tourrent digit_stringg. len [tgta___ggitttssttng)) if boal_digit_string
208
210
-- curtent_digit_ satching_digits
212
clae:
Stop counting as soon as mismatch occurs
break
215
i Scale reward based on pregress
217
if lig.ppalddiit_ssrriaae > 0: biitching_diiittt
0.7
219
220
return min treward, 3.991
Listing I: Android Control "Set Timer" Generated Reward.
24

---

Published as a conference paper at ICLR 2026
A.6.2 MuJoCo: Hopper
imports
import Jax
Import jax.numpy as jnp
# reward Function
fjax-jit
der reward_ In reward_fn :
reward = 0.0
10
Strong bonus for negative values in Teatuzes 0, 1 (expert states are typically negative) state(3]) reward - jap (0 (0.00, state(ll) 2.0 + Honus for negative Seature
reward jap 3np.minimun.00
negative Seature
12
14
Strong penalty for positive values in teatures 0, 1 (learner scates often positive)
15
reward jap m0ximum. 00,0, state(31) 2.0 + Penalty for positive feature 0
16
jnp.maximum state(1])
Penalty positive feature
17
Benus positive values
jexpezt states typicaily positive)
19
reward - jap maximmm state(S]) 3.0 + Bonus for positive feature 5
21
Penalty for negative values in Ieature 5 (Zeazner states often negacive)
22
reward hap.maximum. (0.00 -state(5]) 1.0 Penalty for negative feature
24
Additienal penalties for extreme positive values in key features Festures 2, 5, 10: penslize Large positive values reward - jnp. mam.immm (0.0, state(2]) 5.0
indicator)
26
27
reward jnp. mum (0.0, state[3])
28
reward jnp. maximum (0.0, state(1]) 2.0
29
reward jap maximum m0ximmm(000,, state(6)) -.0
map.maxxium state(7])
31
reward -- jap maximum. state(10]) A 1.0
33
Additiona penalty for large iegative values in features 0,1 (so avoid over-rowarding)
34
reward comp.maximuni (0.0, -state[0]) * 0.5
35
reward jop.maxximum -state(l]) 0.5
36
Listing MuloCo Hopper Generated Reward.
25

---

Published as a conference paper at ICLR 2026
A.7 PROMPTS
Goal Identification Prompt For Data Cleaning
Given this reward code: (reward_code)
Trajectory:
(trajectory)
Please analyze the state sequence and the agent's instruction.
Identify the index of the goal state. The state indices are 1-based.
OUTPUT FORMAT:
Answer in a json format as follows:
"reasoning': : Explain your reasoning for choosing the goal state(s), goal_state_indexes! index of the goal state(s), or if no goal state is present.
list of integers representing the 1-based
Prompt 1: The prompt for identifying the goal state(s) within trajectory.
26

---

Published as a conference paper at ICLR 2026
LLM Initial Reward Generation
You are an Mb engineer writing reward functions for Rl training. Given a trajectory with marked goal states, create Python reward function that can reproduce this behavior.
Requnt:mmnts::
Write self-contained Python 3.9 code
Make the function generic enough to handle variations (different positions, orientations, etc.)
Design for modularity you might extend this reward later to handle multiple goal types
Aim give high rewards for expert states and low rewards for all other states
Environment Details:
(env_codel, (import instruciont (state desare_description
Trajectories
[expert_trajector:i]
Key Instructions:
Analyze the trajectory to understand what constitutes success Identify intermediate progress that should be rewarded
Create utility functions for reusable reward components
The code will be written to file and then imported.
OUTPUT FORMAT:
Answer in a ison format as follows:
'reasoning': Given the reason for your answer
'reward codde'' Code for the Reward function class in the format:
A imports
<imports_here>
# utils functions
<utils functions here>
reward function
class Reward:
def _nnii_ ixtra_info-Nonc)) <code_here>
def reward_fn (self, state): <code_here>
def debug (self, state):
<code here> The Reward class will be initialized with the extra_info argument. Describe in the comments of the class the behaviour you are trying to reproduce.
reward_fn and debug_ receive only state argument. The debug_an should return a string that will be printed and shown to you after calling reward_fn fn on each state. You can print internal class properties to help you debug the function. Extract any needed information from the state store it in the class. The Reward class will re-initialised at the beginning each episode.
Prompt 2: Prompt to generate the initial set of rewards
27

---

Published as a conference paper at ICLR 2026
Evolution Mutation Prompt
You are an ML engineer writing reward functions for RL training. Given a trajectory with marked goal states, create Python reward function that can reproduce this behavior.
Requirements:
Write sel folfconntinedd Python 3.9 code
Aim to give high rewards for expert states and low zewards for all other states
Make the function generic enough handle variations (different pesitions, orientations, etc.)
Design for modularity handle multiple goal types
might extend this reward later to
Original Reward Code: ((code))
([import_mossaeee) ([state__escripption)
CRITICAL: Incorrect Trajectories
The reward function not performing well on the following
trajectories. It either assigned a high zeward to a negative states or assigned low reward to an expert state. The predicted rewards for each step are shown.
Change the reward function fix these errors.
Key Instructions:
Create utility functions reusable reward components
Implement goal switching logic using extra into to determine which reward function to use
Reuse existing utilities where possible
Make sure the logic you write generalises to variations in extra_info
iincoreect.trajectociie
(expert_traj_strr
Now, provide the nutated version of the reward function that addresses these errors.
OUTPUT FORMAT:
Answer in a json format as follows:
'reasening': Briefly explain the corrective change you made. rreaaddcclasscoode
Code for the Reward function class in the
format:
t Reward format and extra info as above
Prompt 3: The prompt used for evolutionary mutation, providing feedback on incorrect trajectories.
28


Reviews:
[{"review_id": "Wp7LZDdIaV", "reviewer": "Reviewer_ZUCH", "summary": "GRACE is an IRL framework that uses code-generating LLMs and evolutionary search to infer executable Python reward functions from expert demonstrations. It operates in three phases: (1) LLM identifies goal states and generates initial rewards from positive/negative trajectories; (2) evolutionary refinement via LLM mutations on misclassified states to maximize fitness; (3) online RL with PPO expands data, enabling further reward improvement. Evaluated on BabyAI and AndroidWorld, GRACE recovers accurate, well-shaped rewards from few demos, outperforms GAIL, matches oracle PPO, and produces reusable multi-task reward APIs.", "questions": "1. How does GRACE perform with open-source code LLMs (e.g., CodeLlama, DeepSeek)? Any ablation on model size or quality?\n2. What is the total LLM inference cost (tokens, time, $) per task? How does it scale with environment complexity?\n3. How reliable is LLM-based goal detection in Phase 1 and 3? Any error rate analysis on false positives/negatives?\n4. Why limit AndroidWorld to Clock app? Have you tested cross-app generalization or real-device deployment?\n5. The method assumes access to environment state (grid, XML)—how does it handle partial observability or raw pixels?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Achieves good reward recovery and generalization with 1–8 expert trajectories, far surpassing neural IRL methods in sample efficiency.\n2. Produces dense, interpretable rewards that enable strong policy learning—matching or exceeding ground-truth PPO in BabyAI and outperforming GAIL.\n3. Demonstrates emergence of modular, reusable reward code across tasks, offering a practical path toward composable reward design.\n4. Clean theoretical link to classical IRL via fitness function, with clear algorithm and reproducible structure.", "weaknesses": "1. Heavy reliance on strong proprietary LLMs (e.g., gpt-4o) without ablation on weaker or open-source models; performance may degrade significantly in practice. The evaluation on other models could be conducted.\n2. No compute cost analysis—evolutionary search with 100+ LLM calls per task is likely expensive and slow, limiting scalability.\n3. Goal state identification by LLM is a critical weak point; no quantification of labeling errors or robustness to hallucinations.\n4. Evaluation scope is narrow: BabyAI is synthetic and low-dimensional; AndroidWorld uses only Clock app with curated negatives—real-world generalization untested.\n5. Baselines (e.g., GAIL) are not compared under identical low-data settings, and multi-task claims lack strong IRL benchmarks.", "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": "zvJJkBrRtY", "reviewer": "Reviewer_2iny", "summary": "Manual reward design for RL is challenging and impractical. To address this, inverse reinforcement learning approaches recovers reward functions from expert demonstrations, but these recovered reward functions are traditionally non-interpretable. Recently, code-generating LLMs recover explainable reward functions but require human-curated task descriptions, goal states or feedback from a trained policy. This paper proposes to recover reward functions purely without task description or goal-specific design specifications by leveraging code-LLMs. First, using the expert demonstrations and random trajectories, initial reward design is created as code. Then, iteratively, through an evolutionary search and a fitness score, reward is refined; using this reward, a policy is trained and additional data is collected, and this process is repeated for fixed number of steps.", "questions": "- Please compare the performance and computational efficiency of the proposed approach with LLM-based reward generation methods in the literature (a,b). \n\n- Please discuss this design selection of using binary reward values in Eq. 3\n\n- Can you please examine the impact of providing environment information to LLM in initial reward design step? \n\n- It is interesting that the goal module is more and more used in later stages of the training, and go_to_reward component is less and less used as training goes on. Just because of curiosity, it would be great to hear the authors' interpretation of this.", "limitations": "", "rating": 2, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "-**Interesting Approach**: Using evolutionary search with a fitness score to generate reward functions looks promising when the size of the expert data is limited. \n\n-**Ablations**: It is clearly presented how the method improves with more expert trajectories and negative trajectories. \n\n-**Analysis of Training**: The proposed approach is well evaluated and how it changes during the training is well presented.", "weaknesses": "- **Computational Overhead**: Computationally the proposed method looks significantly complex. It requires retraining the policy for each reward refinement, and this looks inefficient even with the limited budget $N$  they use for PPO training. The computational overhead of the proposed approach must be clearly evaluated against Adversarial IL (GAIL [c] etc.) and IRL (MaxEnt IRL [d], GCL [e], etc.). Wall clock times during trainings or total number of gradient steps can be compared. \n\n- **Contributions and Novelty**: The main contribution of the paperis claimed to be not using any information other than expert demonstrations. However, in line 161, it is stated that extra information about the environment is used as query to the LLM for initial reward design. Although this is not human-designed, this contradicts the claim of purely using expert data. Please clarify how important this extra environment information is via an experimental analysis. More importantly, the proposed approach is fundamentally very similar to [a], and both of them use the performance of the trained policy as feedback for reward learning. Therefore, the contribution over [a] must be clarified. Please compare both the performance and computational efficiency with [a] extensively.\n\n- **Limited Experimental Setting**: Experimental setting is limited to selected BabyAI environments. Please employ more standardized reinforcement learning and imitation learning benchmarks covering continuous control, locomotion and dexterous manipulation tasks. It would be great if you could evaluate GRACE in the IsaacGym [f] and Bi-dexterous Manipulation [g] benchmarks as in [a].\n\n- **Unclear Design Choices**: In Eq. 3, authors state that they bound the reward function. However, instead of just bounding, they employ binary rewards in Eq. 3 with a threshold. The reason behind this selection is unclear, and not explained. \n\n- **Further Analysis and Clarification**: In Figure 3, the x-axis is unclear, it should be clearly stated whether it is $K$ or not. Also, in addition to the effect o $K$, the effect of $M$ (number of refinement steps) and $N$ (number of PPO training steps *for each refined reward function* should be analyzed. \n\n- **Comparisons with IRL**: It is claimed that the proposed method is more sample efficient than IRL approaches. Please clearly demonstrate empirically that GRACE outperforms IRL (d-e, not adversarial imitation learning which is different than IRL). \n\n[a]: Ma, Yecheng Jason, et al. \"Eureka: Human-Level Reward Design via Coding Large Language Models.\" The Twelfth International Conference on Learning Representations.\n\n[b]: Venuto, David, et al. \"Code as Reward: Empowering Reinforcement Learning with VLMs.\" International Conference on Machine Learning. PMLR, 2024.\n\n[c]: Ho, Jonathan, and Stefano Ermon. \"Generative adversarial imitation learning.\" Advances in neural information processing systems 29 (2016).\n\n[d]: Ziebart, Brian D., et al. \"Maximum entropy inverse reinforcement learning.\" Aaai. Vol. 8. 2008.\n\n[e]: Finn, Chelsea, Sergey Levine, and Pieter Abbeel. \"Guided cost learning: Deep inverse optimal control via policy optimization.\" International conference on machine learning. PMLR, 2016.\n\n[f]: Makoviychuk, Viktor, et al. \"Isaac Gym: High Performance GPU Based Physics Simulation For Robot Learning.\" NeurIPS Datasets and Benchmarks. 2021.\n\n[g]: Chen, Yuanpei, et al. \"Towards human-level bimanual dexterous manipulation with reinforcement learning.\" Advances in Neural Information Processing Systems 35 (2022): 5150-5163.", "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": "CoxHbQw0MM", "reviewer": "Reviewer_JQHA", "summary": "This paper introduces GRACE (Generating Rewards As CodE), a framework that uses code-generating Large Language Models within an evolutionary search procedure to learn interpretable reward functions from expert demonstrations. The method operates in three phases: (1) identifying goal states from expert trajectories, (2) refining reward functions through evolutionary search guided by fitness on goal/non-goal state classification, and (3) expanding the dataset through RL-trained policy rollouts. The authors evaluate GRACE on BabyAI navigation tasks and AndroidWorld device control tasks, demonstrating that it can recover accurate rewards from few demonstrations and outperform GAIL while matching oracle PPO performance. The code-based representation enables interpretability and natural emergence of reusable reward APIs across multiple tasks.", "questions": "- How does GRACE handle continuous states/actions? All experiments use discrete or structured representations. Can it -work with raw continuous sensor data?\n- What is the computational cost (LLM queries, wall-clock time) compared to baselines?\n- Why only GAIL as a baseline? How does GRACE compare to other IL methods?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper addresses an important problem of interpretability in Inverse Reinforcement Learning by representing rewards as executable code rather than opaque neural networks.\n- Using code as rewards is interesting and likely enables faster RL training compared to querying LLM reward models at each step.\n- The approach demonstrates strong sample efficiency, achieving high performance with only 8 expert trajectories compared to GAIL's requirement of 2000 demonstrations.\n- The multi-task experiments on BabyAI showing emergent code reuse and API formation are compelling evidence of the benefits of symbolic reward representations.\n- The writing is generally clear and the method is presented systematically.", "weaknesses": "- Fundamental evaluation concerns: As acknowledged in limitations, the LLM serves as the final judge of success in Phase 3, which compromises evaluation validity since the same model generates and evaluates rewards.\n- Missing critical baselines: The paper lacks comparison to prompted LLMs as reward judges (mentioned in related work but not benchmarked), which would be the most obvious baseline. Similarly, no comparison to using LLMs directly as policies with demonstrations in context.\n- Limited problem scope: The approach reduces reward learning to classifying goal states, which only covers tasks with distinguishable terminal success states. This excludes many RL problems like learning to backflip where initial and final states are identical, or continuous control tasks requiring trajectory-level rewards.\n- Theoretical issues: Proposition 1 is unclear—the relationship between the mask function m and the original IRL objective (Eq. 2) is not explained. How does flipping rewards on goal/non-goal states relate to the max-margin formulation of classical IRL?\n- Statistical rigor: Results appear to be from single runs without error bars or significance tests (e.g., Figure 4 shows single curves). Figure 3 mentions \"3 seeds\" but most other results lack this. Table 1 shows single numbers without variance.\n- Confounded claims: The claim that Phase 3 generates \"well-shaped\" rewards conflates reward shaping with improved coverage beyond expert demonstrations. The performance gains might simply come from seeing more diverse states rather than better shaping.\n- Limited scope: Evaluation restricted to relatively simple domains (BabyAI mazes, basic Android tasks).", "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": "gAf64donJs", "reviewer": "Reviewer_R5gj", "summary": "This paper introduces GRACE, a framework for inferring interpretable reward functions as Python code from expert demonstrations. The approaches iteratively optimizes a population over LLM-generated code-level reward functions via evolutionary search. The initial population is generated from LLM-labelled goal states, and then refined using new samples gathered from an inner RL loop. The data collection and context provided to the LLM in the evolutionary search depend on the goal label provided by an LLM. The proposed method consistently solves BabyAI and AndroidWorld tasks, and outperforms GAIL and a (forward-RL) PPO baseline.", "questions": "I think the present draft is promising, but have some small problems with its rigour and presentation. I’m happy to positively re-assess my rating if the below questions are adequately discussed and addressed.\n\n1. The reliance on LLM-labeled goal states seems to be a major design choice, especially since this labeling is also used for the RL-in-the-loop part of the algorithm. Could the authors provide more justification for this? A discussion on its limitations and potential failure modes, or an ablation showing how performance degrades with less accurate labeling (either via artificial noise or, preferably, from a less powerful LLM), would be very valuable.\n2. The related work section mentions other methods that generate rewards-as-code. Could the authors provide a more detailed distinction between these methods and GRACE and, ideally, add at least one of these methods as a baseline?\n3. Could the authors clarify the experimental setup? How many demonstrations are used for the AndroidWorld tasks, how many seeds are used for the different results, and why?\n4. The notation for demonstrations and their goal state is somewhat overloaded, and a demonstration is essentially only characterized by its goal state. (How) can the method use intermediate states of these trajectories to its advantage? Is it necessary to provide negative trajectories, or can intermediate states act as “negatives” for the LLM prompt? An ablation study\non this behavior would help clarify the capacities and limitations of GRACE", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Using an evolutionary framework where an LLM mutates code-as-a-reward functions based on achieved goal states and failures is a novel and highly promising approach to IRL.\n- By design, the method produces human-readable reward functions, which is a significant advantage over black-box IRL.\n- The paper is generally well-written and easy to follow. The approach is presented in a clear and concise manner, and most hyperparameters and design choices are listed and justified.\n- GRACE shows strong performance on the selected tasks outperforming GAIL with a fraction of the data and producing rewards that are on par with the ground truth in terms of downstream RL performance.", "weaknesses": "- The entire framework seems to depend on an LLM's ability to correctly identify goal states from demonstrations and, more crucially, from new agent trajectories of the inner RL loop. This is a rather strong assumption that deviates from the standard IRL setup. The (potential lack of) robustness of this approach is not analyzed, and it seems to limit GRACE to finite-horizon RL settings.\n- Similarly, the context provided for the LLM-based mutation seems a bit arbitrary and could be better motivated or experimented with.\n- The related work section lists other \"Code as Reward\" methods but does not explain how GRACE compares to or improves upon them. None of these methods are included as baselines in the experiments, and this choice is not justified or explained. This omission makes it difficult to assess GRACE’s position in the research landscape, and consequently the novelty and significance of the method. GAIL is chosen as the only Imitation Learning/IRL baseline, which may be insufficient given a lot of recent advances in, e.g., diffusion-based imitation learning.\n- The experimental results are a bit thin. For example, Figure 3 uses error bars over 3 seeds, which are absent for, e.g., Figure 2. Similarly, Figure 5 reports the mean over three different environments, and so on. Some details, such as the number of demonstrations used for the AndroidWorld tasks, are missing. GRACE solves both presented benchmarks, which makes it hard to judge where the boundaries of the method lie and where it would fail. There is no qualitative visualization, which makes it difficult to appreciate the difficulty of the tasks.", "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)