Trust Region Reward Optimization and Proximal Inverse Reward Optimization Algorithm

We thank Reviewer for the thoughtful and constructive feedback. We appreciate the recognition of PIRO's theoretical contributions, strong empirical performance, and practical scalability. Below, we address the reviewer’s concerns regarding on-policy sample efficiency, policy optimizer stability, and generalization to more complex reward structures.

[Q1,W1,W3]: How does PIRO handle the trade-off between stability and sample efficiency in on-policy rollouts?

Our framework includes strategies specifically designed to address this trade-off. In the original implementation, we used importance sampling (IS) to reuse previously collected rollouts in reward updates. By logging behavior policy probabilities, we reweighted each trajectory’s contribution to the reward gradient. To avoid high variance, we periodically reset the behavior policy using Monte Carlo sampling, which effectively bounds the IS weights and improves sample efficiency without manual clipping.

Under the assumption of sufficiently informative demonstrations, our framework admits an alternative yet theoretically equivalent approach to reward gradient estimation that circumvents the need for computing advantage values $A(s, a) = Q(s, a) - V(s)$ through full trajectory rollouts. Rather than estimating advantages via high-variance Monte Carlo returns, we exploit a closed-form identity that connects the log-likelihood of expert actions with the expected reward difference between expert and learner policies:

This is a known result in offline IRL and non-adversarial IRL (e.g., ValueDICE, SQIL, IQ-Learn), but relying on rich demonstrations. It is theoretically equivalent to maximizing the expected reward difference $\mathbb{E}\_{\pi_E}[r_\theta] - \mathbb{E}\_{\pi}[r_\theta]$, and may offer improved efficiency by avoiding advantage estimation and long-horizon rollouts.

This sampling strategy serves as a practical implementation choice within our broader framework. Empirically, it improves sample efficiency and reduces per-update training cost, while maintaining stable convergence—addressing both the sampling cost and implementation complexity concerns raised in the review.

Table: Training efficiency (time) normalized by PIRO = 1.00×, across MuJoCo and Gym Robotics tasks.

This demonstrates that PIRO's theoretical framework naturally accommodates efficient sampling strategies, reducing dependence on costly on-policy rollouts without compromising training stability.

Regarding the reviewer’s concern on Dependency on On-Policy Rollouts(W1), we agree that standard PIRO updates require on-policy trajectories, which may reduce sample efficiency in interaction-limited settings. However, our framework is flexible and admits off-policy-compatible variants. In particular, the pairwise classification objective introduced above enables fully off-policy reward updates using single-step transitions from a shared reference distribution. This eliminates the need for fresh rollouts during reward optimization and significantly reduces the cost of environment interaction. Our experiments show that this variant achieves comparable performance while improving per-update efficiency, making PIRO more suitable for sample-expensive or offline imitation scenarios.

Table: Mean convergence (the point where significant increases stop and fluctuations happen afterwards) steps across five independent runs on MuJoCo tasks.

Convergence / Rewards

We also take this opportunity to address a related concern from W3 regarding hyperparameter tuning. Specifically, we tested the impact of reward-related coefficients such as $\mu$ and $\nu$ in Appendix H. The sensitivity plots show that PIRO remains stable across a broad range of values, and performance degrades gracefully outside the recommended intervals. This suggests that careful tuning is helpful but not critical, and that the method is reasonably robust for practitioners.

[Q2] What are the implications of PIRO's dependency on a stable policy optimizer in high-dimensional settings?

PIRO is designed as a modular reward learning framework that delegates policy optimization to an external RL algorithm. This dependency enables flexibility rather than limitation: PIRO does not require any specific optimizer and can be paired with both on-policy methods like PPO and off-policy methods like SAC or TD3. In high-dimensional environments, where policy optimization becomes more sensitive to exploration and variance issues, this modularity allows PIRO to leverage state-of-the-art policy optimizers to improve scalability and stability.

As our formulation imposes no structural constraint on the optimizer, the reward learning mechanism remains agnostic to the policy update dynamics, enabling seamless integration with more advanced or domain-specific optimizers. This compatibility is particularly beneficial in high-dimensional settings, where the choice of a robust optimizer can significantly impact performance. In this sense, PIRO's reliance on a stable optimizer is not a weakness, but a strength that supports generalization and extensibility across domains.

[Q3]Can the authors provide insights into the generalizability of PIRO to other types of reward functions beyond state-only rewards?

We discuss two more types of reward functions:

• Action-only reward $r(a)$, which is common in bandit settings. Our method naturally supports this reward setting by setting an action-only reward model $r_\theta(a)$.
• Transition-dependent reward $r(s,a,s')$. In IRL, a common technique to cope with this kind of rewards is to assume a potential-based function $h_\phi(s)$ and use a particular form to model $r(s,a,s')$: $r_{\theta,\phi}(s,a,s') = r_\theta(s,a) + \gamma h_\phi(s') - h_\phi(s).$

Based on this structure, we have $Q_{r_{\theta,\phi}}^{\pi_{\theta,\phi}}(s,a) = \mathbb{E}[r_\theta(s,a) + \gamma h_\phi(s') - h_\phi(s) + \gamma (r_\theta(s',a')+\gamma h_\phi(s'') - h_\phi(s')) + \cdots] =$ $\mathbb{E}\left[\sum_{t=0}^\infty \gamma^t r_\theta(s,a) | s_0=s, a_0=s\right]$

Thus, the Q-function gradient under this shaped reward formulation is mathematically equivalent to that used in our original derivation, resulting in complete compatibility with the PIRO update mechanism.

We sincerely thank the reviewers for their constructive and thoughtful feedback. We are encouraged that the reviewers found our theoretical contribution principled and insightful, appreciated the rigorous derivations, and recognized the breadth of our empirical evaluations across both simulated and real-world domains. Below, we address reviewer’s comments in turn.

[W1,W2,Q1]: In practice, does using heuristic approximations (e.g., adaptive coefficient $C$) break the theoretical monotonicity guarantees of PIRO? Can you quantify how often KL divergence increases during training?

Indeed, using a relaxed adaptive coefficient $C$ cannot rigorously guarantee the monotonicity in theory. Theoretically, with an adaptive coefficient and assuming we have a space of reward parameter that makes the KL divergence Lipschitz continuous, we can then quantify that KL divergence increase is upper bounded by $\mathcal{O}(\frac{C\epsilon}{1-\gamma})$. This is because $Q$ is $\frac{1}{1-\gamma}$-Lipschitz in $r$ and the energy-based policy $\pi$ is $1$-Lipschitz in $Q$ (assuming the temperature = 1). Empirically, we found that the rate of KL divergence increases is lower than baselines in general. On the other hand, approximating the theory for better feasibility and efficiency is common in reinforcement learning approaches. Closely related to our method is TRPO, which also employs a heuristic KL constraint to replace the theory-informed one that guarantees the monotonic policy improvement. At its reputation-gain successor PPO, the KL constraint is further simplified as a penalty term with an adaptive coefficient (or clipping). This simplification makes the great success of PPO, which is now the gold-standard of RL algorithms. In light of this, we believe our methods -- inspired by theory but not rigorously following the theory -- are in line with the mainstream design and practice of RL algorithms. It is also worth noting that although some prior methods also adopt the concept of policy alignment, no one, at least in theory, analyzes how to achieve monotonic alignment, but we did in this work. This sets our work apart from the literature.

[Q2,W3,W4]: Can you elaborate on when the trade-off between stability and sample efficiency is acceptable, and whether off-policy correction techniques (e.g., truncated IS) could be used to alleviate this?

This trade-off is acceptable in settings where interaction is cheap (e.g., simulated control), or where expert trajectories are fixed and reward learning is purely offline (e.g., behavioral modeling). For instance, in MuJoCo, stable reward updates are more beneficial than marginal gains from reuse. In real-world applications like sim2real robotics or clinical modeling from logs, safety and interpretability matter more than sample count—making stability a higher priority than efficiency.

By contrast, in real-time robotics or interactive recommendation systems—where every environment step is expensive or delayed—this trade-off becomes limiting, and improving sample efficiency becomes essential.

As for off-policy correction, we already employ importance sampling in our reward updates, using behavior-policy log-probabilities stored in the trajectory buffer. Instead of explicit truncation, we control variance via periodic Monte Carlo resampling, which resets the behavior policy and keeps IS weights bounded. This approach achieves similar stability benefits without requiring manual clipping thresholds.

In addressing this question, we also aim to clarify two concerns raised in the weaknesses: the reliance on on-policy sampling (on-policy cost) and the complexity of gradient estimation .

Within our framework, when demonstrations are rich, we can adopt an equivalent but potentially more efficient formulation of reward gradient estimation that eliminates the need for estimating $A(s) = Q(s,a) - V(s)$ on trajectory rollouts. Instead of relying on advantage estimates derived from sampled returns—which often involve high-variance Monte Carlo evaluations—we directly optimize the expected reward difference between expert and policy transitions according to the following equivalence relationship:

Here, $d^\pi_t(s)$ denotes the state distribution induced by $\pi$ at $t$ and $\eta$ is the (known) initial state distribution. This is a known result in offline IRL and non-adversarial IRL (e.g., ValueDICE, SQIL, IQ-Learn), but relying on rich demonstrations. It is theoretically equivalent to maximizing the expected reward difference $\mathbb{E}\_{\pi_E}[r_\theta] - \mathbb{E}\_{\pi}[r_\theta]$, and may offer improved efficiency by avoiding advantage estimation and long-horizon rollouts.

This sampling strategy serves as a practical implementation choice within our broader framework. Empirically, it improves sample efficiency and reduces per-update training cost, while maintaining stable convergence (see Table below). This addresses both the sampling cost and implementation complexity concerns raised in the review.

To further demonstrate PIRO’s sample efficiency, we report the number of iterations required for each algorithm to reach convergence across a range of MuJoCo tasks. The table below summarizes the average convergence steps of five independent runs, providing a direct comparison of training speed under consistent evaluation protocols.

Table: Mean convergence (the point where significant increases stop and (small or big) fluctuations happen afterwards) steps across five independent runs on MuJoCo tasks. Format: Convergence / Rewards

These results demonstrate that although PIRO may not exhibit the fastest raw convergence in terms of iteration count, it consistently achieves stable and reliable progress across tasks. Unlike methods that converge quickly but suffer from regressions or collapse, PIRO guarantees monotonic policy improvement (in theory)—ensuring that each update brings the learner closer to the expert. As a result, PIRO avoids unnecessary detours and converges toward expert behavior in a more direct and (empirically) sample-efficient manner. This highlights the practical benefit of our framework: trading off aggressive updates for principled, steady advancement leads to faster and more stable convergence overall.

[Q3]: Can you compare the core strengths and limitations of PIRO’s monotonic policy alignment versus value-matching non-adversarial methods like IQ-Learn?

Both PIRO and IQ-Learn avoid adversarial training by replacing discriminator-style reward objectives with more stable surrogate losses. The key difference lies in what is aligned with the expert: IQ-Learn focuses on value matching—training a value function such that its gradient implicitly guides policy learning—whereas PIRO explicitly ensures monotonic reduction in policy divergence from the expert via a proximal reward update mechanism.

The main advantage of PIRO is that it directly optimizes for behavioral alignment between the learned and expert policies. This leads to predictable, stable policy improvement steps with a clear theoretical guarantee (via majorization-minimization). In contrast, IQ-Learn's value-matching objective does not guarantee that policy divergence decreases, and its performance can depend heavily on the quality of value estimates, which are sensitive to bootstrapping errors.

That said, IQ-Learn can be more sample-efficient in certain settings due to its compatibility with off-policy value-based learning, while PIRO is more conservative by design. Our approach prioritizes stability and monotonic policy improvement, which is particularly beneficial in settings where safe or interpretable reward learning is critical.

Dear reviewer,

The authors have posted a detailed response. Does it assuage your concerns? I struggle to find a concrete reason for rejection from your review. I am a bit at a loss when you say that on one hand, it's insightful theory, and on the other that some things are not theoretically justified.

Thanks for your engagement and comments. Could you please specify which parts of your concerns have been addressed and which part not? Also, could you please take a look at and reply to AC’s questions above (ignore this if you have already replied to it and keep it confidential for the authors)? Your detailed information will be helpful to us. Thanks in advance.

We thank the reviewer for the thoughtful feedback. We're glad you found the "presentation extremely clear", the diagrams "fantastic", and the method "intuitive". We're especially encouraged by your belief that the "main algorithmic idea will be of broader interest to the community". We address your questions below.

[Q1]: Did you try directly minimizing the KL to the expert, and was it unstable?

Yes. We experimented with directly minimizing the policy KL divergence $\mathrm{KL}(\pi_E | \pi_\theta)$ via gradient descent. However, as shown in the table below, this naive approach exhibits significant instability across tasks. In environments like CartPole and Q*bert, performance dropped sharply and fluctuated erratically. Even in more stable tasks like Humanoid, it failed to consistently improve the policy.

These results reflect common issues with directly optimizing KL: the gradients can be high-variance and poorly aligned with actual policy improvement, especially when $\pi_\theta$ is far from $\pi_E$. This was a key motivation behind our design of PIRO. By using a majorization-minimization (MM) framework to upper-bound the KL surrogate at each iteration, PIRO ensures stable updates with guaranteed monotonic improvement—effectively avoiding the oscillations and regressions observed with naive KL minimization.

To illustrate this behavior concisely, we select one representative environment from each category (Gym, MuJoCo, Atari, Box2D), due to words constraints.

Table: Return comparison between PIRO and naive KL minimization over iterations. Directly optimizing KL divergence results in instability (high fluctuation).

[Q2-1]: How sensitive is PIRO to the inner policy optimizer? Did you try other methods beside PPO?

PIRO’s core theoretical framework does not rely on a specific policy optimizer. It only assumes access to a differentiable policy that supports sampling and advantage estimation. Therefore, in principle, any optimizer—on-policy or off-policy—can be used, as long as it enables stable policy updates and maintains sufficient overlap with expert trajectories.

In response to this question, we conducted additional experiments using Soft Actor-Critic (SAC) in place of PPO for policy optimization.

Table: Performance across 5 independent runs on MuJoCo control tasks using high-quality expert demonstrations trained with SAC.

Table : Performance across 5 independent runs on Gym Robotics tasks.

These results demonstrate that PIRO is not sensitive to the choice of policy optimizer, and performs robustly under both PPO and SAC when compatible conditions are met.

[Q2-2] Is it correct that any method which learns the EBM (and can sample from it) would work

In theory, any method which learns the EBM (and can sample from it) is compatible with our framework. But this is not to say that they will work well because the sampling mechanism should be well designed to make it able to cope with high-dimensional and continuous states or actions. Given that, SAC would be a suitable alternative policy optimizer and its peformacme can be found above.

[Q3]: Is there any kind of clipping used to control the importance ratios in the gradient updates? Empirically, do these not blow up?

We do not apply explicit clipping to the importance ratios in PIRO. Instead, we reset the behavior policy used for advantage computation to the current policy after a small fixed interval. This guarantees the current policy will not deviate from the old policy too much, ensuring that importance weights remain well-behaved.

Furthermore, our implementation buffers and reuses trajectories based on specific (state, action) pairs. This design amortizes advantage estimation across updates and avoids repeatedly sampling from stale or highly divergent policies. Empirically, this mechanism helps contain the variance of importance weights. Across all tasks, we observe no gradient explosions related to importance sampling—learning.

[Q4]: For Q and V gradient updates, could you maintain a buffer of all off-policy rollouts and probabilities seen up to that point? Do you have thoughts on whether PIRO would stably support this kind of change to improve sample efficiency?

While PIRO is primarily on-policy in its reward update logic, our implementation does maintain trajectory buffers for off-policy reuse—specifically for estimating Q and V values under the old policy and current reward. These buffers are keyed by individual $(s, a)$ or $s$ pairs and store full rollouts sampled from the behavior policy. The associated log-probabilities from the behavior policy are retained to enable importance sampling (IS) during advantage estimation.

This design is crucial for our IS-based reward updates, as reusing cached rollouts allows us to amortize gradient estimates without repeatedly querying the environment. To avoid accumulating divergence, we do not maintain an ever-growing replay buffer. Instead, we periodically clear and refresh the buffers via Monte Carlo re-sampling (see Table 2), which resets the behavior policy to the current policy snapshot. This bounds policy shift and ensures that the IS weights remain well-behaved, as discussed in [Q3].

[Q5]: why on-policy robustness of PIRO seems to differ per task?

While PIRO theoretically ensures monotonic alignment at the reward level, its practical robustness depends on how reliably policy updates respond to reward signals. Hopper, despite having nonlinear dynamics, tends to exhibit smooth and dense reward landscapes and stable PPO training behavior. This makes it less sensitive to minor variations in reward updates, hence showing strong performance and stability under PIRO.

In contrast, tasks like Ant or HalfCheetah often involve higher-dimensional control and more complex reward trade-offs (e.g., speed vs stability), where small reward shifts can induce divergent behaviors. Since PIRO does not directly constrain policy updates but rather modulates them via learned rewards, the resulting robustness can vary across environments depending on policy sensitivity to the shaped rewards.

[Q6]: For reward recovery in Appendix E, does the PIRO-learned policy match the expert performance?

Yes. The experiment in Appendix E is conducted in a 6×6 GridWorld environment with a manually designed ground-truth reward that depends only on the agent’s position. The PIRO agent is trained without access to the true reward and learns solely from expert demonstrations. The reward heatmaps confirm that the recovered reward closely matches the ground-truth.

To verify behavioral alignment, we measured the average return of the learned policy over the last 5 evaluation rounds and compared it to the expert return:

This indicates that the PIRO-learned policy not only matches but slightly exceeds the expert's performance in terms of accumulated reward.

[Q7]: Why is it reasonable to assume a shared policy across individual meerkats in the dataset? The meerkat behavivor dataset only records meerkats' physical behavior, like transiting among locations in their habitat, foraging and play-fight. These physical behaviors do not include social, sychological or ethical behavior, and thus they are largely same among adult meerkats. This assertion is supported by the records in the dataset, where each individual meerkat has a unique track number in a video clip, one can find no outstanding differnece over behavior of different individual meerkats. For example, raising guard would be a supporting evidence, as different adults may act as guards in turn.

1. How is the reset frequency determined

We tested the policy reset frequency within the range [5, 10, 20, 50, 100]. A large frequency reduces the computational burden but increases the variance, while a too too samll frequency improves accuracy but raises the computational load. We thus choose 10 which demonstrates a tradeoff betweem accuracy and computational efficiency.

2. Stability of Importance Sampling Weights

We didn't apply clipping to the importance ratios in our experiments. However, we find that the product of importance ratios do not blow up or shrink in general; see the Table below for our reported results of importance ratios in Cartpole-v1. We owe this stability to the proximal reward update, which leads to samll policy changes and thus stable importance ratios. As you suggested, we also add the clipping operation to the importance sampling weights. The results indicate that the clipping could reduce variances and improve reward performance, though not that significantly (see Table below). This suggests us to incorporate clipping into our final released code for public use.

Table: Important sampling weight comparison in Cartpole-v1.

Note 1: In PIRO, each global roound process several batches, each with an importance sampling weight. We report the maximum, mean, and standard deviation of the importance sampling weights among these batches.

• Reset Frequency = 10.

Note 2: As shown, clipping helps stabilize the mportance sampling weights early in training (Rounds 2–8), while maintaining slightly higer or comparable rewards to the unclipped case. After the first reset (Round 10), both configurations converge to similar mportance sampling weight (around 1) with low variance and near-unity mean weights, demonstrating the effectiveness of proximal reward update that leads to small and stable policy change.

3. Re: Sample efficiency

Thanks for pointing our this. We have made a clarification as follows. The sample efficiency that Reviwer zT3i mentioned is wrt online sample efficiency, while we cliamed "sample efficient" is wrt "expert demonstration samples", i.e., PIRO could demonstrate good performance with small amount (one trajectory) of expert samples, as verified in our previous responses. I deed, as you pointed out, our mian contribution is on the stability while not sacrificing the final performance of reward recovery and policy imitation. We have acknowledged the limitation in the experiment and conclusion section that: in general, the stability comes at the cost of more frequent online sampling . We will continue to explore fully offline IRL with stability guarantees, which we believe will further contribute to the reinforcement learning and imitation learning communities.

We thank the reviewer for the helpful and constructive feedback. We are glad they found our approach a “novel look at proximal reward updates” and appreciated the "beautiful and informative figures”.Below we provide our responses, and we’re happy to follow up on any additional questions.

[Q1,W1,W2] How sample-efficient is PIRO compared to SOTA off-policy methods, such as DAC, P2IL and CSIL?

As correctly noted in W1, PIRO is an on-policy method, and our evaluation thus focused on comparisons with other on-policy baselines. Meanwhile, to address the concern in W2, we thank the reviewer for pointing out our wording issue. While PPO was used to train expert policies across most of the baselines (for consistency), IQ-Learn was implemented with SAC as its policy optimizer, as in the original paper. We will revise the manuscript to clarify this distinction.

Given the reviewer’s valuable suggestion, we conducted additional experiments to evaluate PIRO’s sample efficiency against DAC, P2IL, and CSIL. All methods use the same high-performing expert models trained with SAC to ensure fair comparison. We will present more detailed results in our response to Q2.

Following standard IRL benchmarking protocols, we assess sample efficiency by measuring the final performance achieved within a fixed number of environment steps (1M, relatively small in the practice of IRL approaches). Results (averaged over three runs) on representative MuJoCo tasks are presented below:

Relative Computation Cost per Batch (Normalized)

Table: Average per-batch computation time normalized by PIRO = 1.00×. Measured over the same training rounds and batch structure across methods.

These results demonstrate that PIRO consistently achieves competitive or superior final performance across MuJoCo tasks. Notably, while DAC performs well in Hopper and HalfCheetah, PIRO shows substantial gains in more challenging environments like Ant and Walker2d, indicating overall better sample efficiency across diverse settings. Although PIRO incurs slightly higher per-batch computation time in some tasks, we observe that DAC is often slower overall. Importantly, PIRO’s optimization framework guarantees monotonic policy improvement at each reward update, enabling more stable training and faster convergence in practice. This tradeoff between efficiency and stability ultimately leads to superior overall performance.

[Q2,W4] The results in Figure 3 look nowhere near optimal

The key reason is that we use PPO-trained experts to match the learner’s optimization and ensure consistency. Prior benchmarks (e.g., Tianshou, SpinningUp) have consistently shown that PPO-trained experts tend to achieve lower episodic rewards than SAC-trained ones, particularly in MuJoCo tasks such as Hopper, Walker, and Ant.

Moreover, we use the latest version of the MuJoCo physics engine (v4), which adopts stricter physical dynamics and reward definitions. This results in significantly lower raw reward values compared to older versions of the environment, even for well-trained expert policies.

As a result, the reward ceilings in our Figure 3 plots are naturally lower—not because PIRO performs poorly, but because the reward scale is inherently limited by the PPO expert ceiling.

This also addresses the concern in W4 regarding the missing expert reference: we now report expert performance explicitly. Below we report the approximate expert rewards and standard deviations used in our study:

In light of the reviewer's insightful comments, we recognize the limitations of relying solely on PPO-trained experts and on-policy pipelines in complex environments. As a result, we have extended PIRO to support SAC-based policy optimization and conducted additional experiments on representative MuJoCo locomotion and Gym Robotics tasks.

Below are the updated reward metrics averaged over five independent runs:

Table: Performance across 5 independent runs on MuJoCo control tasks using high-quality expert demonstrations trained with SAC.

Table : Performance across 5 independent runs on Gym Robotics tasks.

[Q3,W3]: On early convergence and update granularity

As noted in our response to Q2, we use PPO-trained expert policies. These tend to converge quickly but may not reach the highest possible reward ceilings achievable by more exploratory algorithms such as SAC. As a result, PIRO often aligns with these experts in just a few global rounds, leading to flattened performance curves after early updates.

Importantly, this behavior reflects PIRO’s high sample efficiency rather than stagnation. Unlike many baselines that may require extended exploration or suffer from unstable updates, PIRO benefits from a theoretically grounded guarantee of monotonic policy improvement. Each reward update provably brings the learner policy closer to the expert, avoiding regressions and unnecessary detours. As a result, PIRO not only ensures stable convergence but also reaches expert-level performance more directly and with fewer interactions—especially in simpler environments such as CartPole, where optimal behavior is quickly attainable.

We report learning progress using global rounds, where each round comprises multiple environment steps followed by one full policy and reward update. This unit more accurately reflects algorithmic progress than raw environment steps.This may also clarify the concern raised in W3: although update-based metrics are implementation-dependent, all baselines in our study are evaluated under the same global-round setup with consistent hyperparameters. None of the results are borrowed from prior work — all experiments are conducted within our unified evaluation framework.

In short, PIRO achieves monotonic improvement under this regime while also demonstrating faster convergence than many baselines in practice.

W[5]: Missing references The comparison with CSIL can be found above. In practice, BC-IRL can be seen as a special variant of ours by removing constraint on the reward update magnititute, except for the way of calculating loss function: BC-IRL minimizes the mean squared loss $\mathbb{E}_{(s,a)\sim \pi_E}(\pi_\theta(s) - a)^2$ while our method minimizes the KL divergence. Least-squares inverse Q-learning (LSIQ) penalizes the reward function magnitude and gives the theoretical support for doing so; we also make a similar penalization through a regularization term (see Eq.(10) in Section 5.1).

Our method stands out as all these prior methods take no constraint on the reward update magnitude, making them vulnerable to training instability. The performance degradation caused by removing the control over reward update magnitude can be found in our ablation study in Appendix G.

Thank you for your continued engagement and detailed feedback. We address your three main concerns below:

TL;DR (Summary of this comment box)

All new baseline implementations (DAC, CSIL, P2IL) are based on their official repositories without hyperparameter modifications. In our paper, for experimental consistency, we trained all expert policies using PPO. This naturally results in lower absolute reward values for both experts and learned policies compared to studies using SAC experts. However, as demonstrated in our supplementary experiments with SAC-based experts, PIRO continues to achieve superior performance relative to other methods. The fair comparison under identical experimental conditions validates PIRO's effectiveness, and it remains robust (stable training) regardless of the policy optimizers (PPO or SAC). The choice of different experts doesn't hurt our main contribution: the first theoretcal framework for principled stable IRL and an algorithm that transforms the theory into practice.

1. Suboptimality Issue

You are correct that our reported results are 5–10x smaller than those reported by Orsini et al. This is primarily due to our choice of expert training algorithm.

Orsini et al. mentioned:"For the Gym tasks, we generate demo with a SAC agent trained on the env reward"

As we mentioned in our Q1 response, we used PPO experts for consistency across all methods. PPO inherently achieves lower performance compared to SAC in continuous control tasks. This is well-documented in established benchmarks. For reference, here are the Tianshou benchmark results comparing SAC vs. PPO:

We could not find publicly available better PPO expert models for the v4 environments during our experiments, so we trained the experts ourselves. Our PPO expert training used the following consistent hyperparam. across all envs:

Only the total training steps were adjusted for different envs based on their complexity:

As we mentioned in our response to Q1, the performance of our experts has already been provided. We also appreciate your insightful question, which prompted us to conduct additional experiments using newly trained SAC experts, closely aligned with those used in Orsini et al. The results of these new experiments have been included in our response to Q2.

2. Baseline Implementation and Results

We implemented all baselines using their original source code and hyperparam. to ensure fair comparison.

With a few noted exceptions:

• Implementation: Most baselines were implemented using the official source code released by the original authors. The only exceptions were GAIL, AIRL & BC, which were implemented using the imitation library.
• Expert demonstrations: A single expert trajectory containing 1000 state-action pairs was used for all methods. (some original papers report results using multiple expert traj.)
• Expert policy: The same SAC-trained expert was used consistently across all baselines.

Some baselines do perform well in specific tasks (e.g., DAC in HalfCheetah, GAIL in Hopper & ML-IRL), which aligns with their strengths. However, PIRO demonstrates more stable training process across diverse envs due to the principle constraint on reward update magnititude.

3. MuJoCo Version Performance Analysis

Our claim about v4 having stricter dynamics was based on our empirical observations during development. To provide concrete evidence, we conducted a comparison of expert performance across MuJoCo versions using same PPO training parameters :

The following table reports the PPO expert performance on MuJoCo-v2 and MuJoCo-v4 envs:

Empirically, we observe that expert performance in MuJoCo-v4 environments tends to be slightly lower than their v2 counterparts in most tasks. This supports our earlier intuition that v4 introduces subtle changes that may affect policy performance.

Different training configurations—including varying training steps, learning rates, and other parameters—can cause substantial performance variations, so one may encounter higher performance results on Hugging Face that used different experimental setups.

Thank you for the comments. We reproduced all baselines using their official code and default hyperparameters with identical expert demonstrations for fairness. We did not fine-tune any method. As mentioned previously, we use a single expert trajectory across all methods to ensure experimental consistency. This also explains why several off-policy algorithms underperform, as they are more sensitive to the number of provided expert demonstrations. This explains why our reproduction does not match the high performance in the original paper. It is not due to a lack of diligence in reproducing prior results, but rather our commitment to ensuring experimental fairness.

Regarding the W&B Results

We had observed the higher performance numbers from the W&B project you referenced earlier in our experiments. However, in order to ensure fairness and consistency, we chose to rerun by ourselves using identical expert demonstrations and training protocols.

We believe the main cause of this discrepancy lies in the limited number of demonstrations: As we stated previously, all methods are trained on the same fixed demonstration to ensure consistency. However, this design also limits demonstration diversity, which significantly impacts off-policy methods. We kindly refer the reviewer to Figure 6 in CSIL (Watson et al.), which shows that CSIL, IQ-Learn, and PPIL all perform poorly in the offline setting when trained on a single expert trajectory — results in that are fully aligned with our findings.

Table: Estimated Normalized Returns from Figure 6 (Watson et al., CSIL) Approximate normalized return values(by expert) under offline imitation learning using 1, 3, 10, and 30 demonstrations across four Gym MuJoCo tasks.

Different expert samples from the same expert policy may further affect outcomes as well. We believe the limited trajectory count is the dominant factor, which, together with the effect of different expert demonstrations, causes the performance differences between the CSIL performance reported in the original paper and in our experiments.

These performance differences do not necessarily indicate poor algorithmic design, but instead, the results reflect that PIRO demonstrates superior stability, consistently achieving steady reward improvement under the same expert conditions where other methods may show greater instability.

We thank the reviewer for the enthusiastic participation in this discussion.

Summary

It seems that this discussion has fallen into a complex cycle. Let's make a long story short through a new experiment that addresses your concerns all at once.

In this new experiment, we make the following settings:

• PyTorch re-implementation of CSIL (Online version). After several attmepts, we still cannot reproduce the performance of the original Tensorflow-based repo of CSIL as reported in its paper (Watson et al.). We thus re-implement CSIL based on PyTorch.
• Fine tune CSIL to match the expert reward performance.
• Env steps as the stastistical unit. This responses to your cocnern on not using env steps.
• Reward performance and stability as performance indicators.
• We use a single demonstarted trajectory (1000 state-action pair).

Experiment results (TL;DR): Both CSIL and our PIRO can reach expert's reward value with PIRO slightly outperforming CSIL, but our PIRO shows better learning stability. See the tables below for the detailed results. This is in line with our theoretical contribution: a stable IRL framework.

Discussion (TL;DR): We appreciate the good performance of CSIL on rewards. Our main contribution is a theory-grounded IRL framework which can maintain high reward performance. Do you agree with this?

Experiment Results

Table: PIRO vs. CSIL Performance.

• Format: Rewards (Increase Rate (Δ%)). A lower average absolute Increase Rate (|Δ|%) indicates a more stable training process.*

Table: CSIL Key Parameters

1. Sample efficiency.

• In our previous response, we used our own expert demonstrations, not those from Orsini et al. or Watson et al. As such, the expert rewards we reported may differ from theirs, as previously noted.

• We are additionally carrying out experiments using expert demonstrations from Orsini et al. Results for CSIL and our PIRO will be coming soon, and we will report them in a new author comment.

• Regarding your question on sample efficiency, thanks so much for your time in collecting and reporting this valuable data! This data verifies the fact that baselines are vulnerable to instability. The sample efficiency you mentioned is, more precisely, sample efficiency wrt online sampling. In our submission, we have acknowledged that our method improves the stability at the cost of more frequent online sampling -- as a result of conservative updates -- in general (e.g., on some MuJoCo tasks). The “sample efficient” we claim in the submission is wrt “expert demonstration samples”, i.e., PIRO could demonstrate good performance with small amount (one trajectory) of expert samples, as verified in our previous responses. We want to clarify this point.

Indeed, in modern RL (recall TRPO vs. standard policy gradient methods), proximal updates will lead to small updates and more frequent sampling operations. In the inverse problem space, our PIRO vs. IRL baselines can be understood as the counterpart of TRPO vs. standard policy gradient methods.

When concerning about "SOTA", based on our results and your results, we could claim that our PIRO is SOTA in stability and final reward performance (policy imitation), though not in sample efficiency wrt online sampling.

Designing an IRL algorithm that guarantees stability and high online sample efficiency won't be easy. We will continue the exploration toward it. We will also explore fully offline IRL with stability guarantees, which we believe will further contribute to the reinforcement learning and imitation learning communities.

2. Joe Watson's official repository

Sorry for the confusion. Joe Watson's official repository is, indeed, based on JAX. It incorporates TensorFlow utilities for expert data generation, processing, and some other computations.

3. Additional discussion on IQLearn

Regarding the LSIQ paper by Al-Hafez et al. (2023) you mentioned, we took a closer look at their reported results. While they show IQLearn reaching 80% on Hopper at around 250K steps, the reward actually drops or fluctuates significantly after reaching an early peak, suggesting that the policy imitation is not that stable. On Ant and Humanoid, our reproduced IQ-Learn results reach the maximum reward with fewer environment steps than reported in LS-IQ. On Walker2d, however, IQ-Learn in our reproduction takes more steps. For reference, we also list the IQ-Learn results from the paper of Hybrid IRL (Ren et al., 2024)

Table: Reproduced Results of IQ-Learn — Environment Steps to reach maximum reward (PIRO vs. LS-IQ vs. HyPE)

We have finished new experiments using the expert demonstrations from Orsini et al., and we indeed reproduced the strong performance of CSIL under this setting. CSIL converges quickly and achieves quite high returns which shows impressive online sample efficiency.

However, as you also observed, CSIL may show larger fluctuations during training, while PIRO maintains a much smoother learning curve. This highlights a key strength of our approach: training stability.

Re: convergence speed. PIRO does not use behavioral cloning (BC) pretraining, which may explain why it does not perform as well as CSIL during the initial training phase.

To summarize, while CSIL is more sample-efficient in online sampling, PIRO achieves comparable or superior final performance and training stability across most environments (Hopper, Halfcheetah and Ant). Based on these results, we believe PIRO is SOTA in terms of stability and final reward performance, though not in online sample efficiency (which is not our intention when developing the theory).

PIRO vs CSIL Performance Comparison (Normalized)

HalfCheetah (expert reward：8770)

Hopper (expert reward：2798)

Walker2d (expert reward：4118)

Ant (expert reward：5637)

Final Performance Summary

Thanks for comments.

Throughout, we make no claim that our contribution is on "online sample efficiency" but the theoretical and practical stability.

And we believe one should not evaluate a method solely based on the SOTA performace on all aspects (accuracy, efficiency & stability) but also based on the methodological innovation. Critically, our method can beat CSIL, IQ and other online baslines on final reward peformance (accuracy).

We have to wrap up the discussion here as the deadline of author-reviewer discussion approaches. Thanks for your enthusiastic participation again. Your suggestions are of no doubt invaluable to us.

In fact, we hypothesize that ”accuracy — online efficiency— stability” is a hard-to-achieve or even impossible triangle for online IRL methods. This work prioritizes the edge between accuracy and stability in this triangle while keeps online efficiency at an acceptable level.

                                    stability
                                       / \
         our prioritization   <----   /   \
                                     /     \
                     reward accuracy ------- online efficiency