Learning Transparent Reward Models via Unsupervised Feature Selection

Authors: daulet  baimukashev
gokhan  alcan
kevin sebastian luck
ville  kyrki
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-07 13:45:04

Abstract: In complex real-world tasks such as robotic manipulation and autonomous driving, collecting expert demonstrations is often more straightforward than specifying precise learning objectives and task descriptions. Learning from expert data can be achieved through behavioral cloning or by learning a reward function, i.e., inverse reinforcement learning. The latter allows for training with additional data outside the training distribution, guided by the inferred reward function. We propose a novel approach to construct compact and interpretable reward models from automatically selected state features. These inferred rewards have an explicit form and enable the learning of policies that closely match expert behavior by training standard reinforcement learning algorithms from scratch. We validate our method's performance in various robotic environments with continuous and high-dimensional state spaces.

Main paper: 
 introduction: 
 Imitation learning (IL) involves learning policies from expert demonstrations. It has been successfully applied in many real-world robotics applications such as household tasks [1, 2] , manipulation [3, 4, 5] , and autonomous driving [6, 7] . It is more straightforward for humans to demonstrate task execution than to precisely specify task descriptions and/or learning objectives [8] . More broadly, IL can be addressed by either behavioral cloning [9] or inverse reinforcement learning (IRL) [10] . The first approach learns the mapping from state to action using supervised learning but suffers from compounding errors outside of the training distribution [11] . Reinforcement learning (RL) [12] benefits from exploration of the environment and learns through experience. RL has been applied in robotic applications [13, 14, 15, 16, 17] , but it requires a well-defined reward function to solve the task. Manually specifying a reward function poses additional challenges for complex tasks, creating a need for the automatic construction of rewards from data. We propose to construct reward by automatically selecting its component features and retrieving reward using inverse reinforcement learning (IRL) setting [18, 10] . IRL benefits from expert demonstrations while utilizing more exploration.One of the prominent approaches for learning rewards is maximum entropy inverse reinforcement learning (Max-ent IRL), which addresses the reward ambiguity problem using a probabilistic model of the behavior [19, 20, 21] . Max-ent IRL can learn an explicit reward model, but typically the features that compose the reward are specified beforehand. Both linear [22, 23, 24] and non-linear [25, 26] functions of features have been used to learn the reward. As depicted in Fig. 1 our method does not rely on hand-picked features as in current IRL methods, but finds relevant set of features from large candidate set. We learn transparent reward model such that the relationship between Figure 1 : Current IRL methods require hand-picked features to use as reward components. Adversarial methods learn the reward using neural networks which are less transparent and not amendable. But, the proposed method finds automatically relevant set of state features and constructs transparent reward models. state features and the resulting reward is direct, and reward function is compact and amendable if preferred.In this work, we propose to construct a reward function from state features using unsupervised feature selection without requiring access to ground truth rewards. To achieve this, we use the maximum entropy formulation for stochastic policies, in which the probability of a trajectory is proportional to its cumulative reward. This approach mitigates the need for an exhaustive search for suitable features and produces an explicitly represented reward model. As only state features are used for reward components, this reward formulation is also referred to as disentangled rewards and is robust to changes in the environment dynamics, as presented in [27] . Our contributions are as follows:1. We develop a method for learning a compact and transparent reward function for use in reinforcement learning algorithms.2. We demonstrate the effectiveness of our proposed method through validation on various continuous control systems with high-dimensional, continuous state spaces.2 Related Works 
feature construction for irl: 
 Inverse reinforcement learning methods can retrieve rewards if their structure is specified accordingly. Many works have considered manually designed features [28, 23, 20, 29, 30] for reward. While human-specified rewards are effective in some scenarios, given only demonstrations, it is challenging to design a reward function that matches the true intent or reward of an expert. The literature on automatically learning compact reward structures is limited. Reward construction from atomic features using regression trees was proposed in [31] . This method iteratively constructs the features and corresponding rewards. Another work [26] considered Gaussian processes for modeling nonlinear function approximation for the reward. A Bayesian approach using similar atomic features that constructs logical conjunctions of features was presented in [32] . Authors in [33] pre-train feature vector using self-supervised loss, but uses trajectory rankings to learn the reward function. Compared to previous approaches, we do not assume to have relevant atomic features or preference labels. Instead, we use states only from expert trajectories and consider nonlinear basis functions of states as reward components. This results in an explicitly represented reward model that consists of a small number of features.An alternative approach for imitation learning, effective in high-dimensional tasks, is enhanced with adversarial methods as presented in [27, 34, 35, 36] . These methods adopt the generative adversarial networks (GANs) [37] training scheme to generate trajectories similar to those of experts. The cost of the discriminator is usually referred to as a learned cost or reward. While adversarial methods achieve state-of-the-art performance for a wide range of applications [38, 39, 40] , the retrieved reward function is a neural network and not an explicitly represented reward model. 
problem statement: 
 Reinforcement learning (RL) is an approach to solve Markov Decision Processes (MDP) defined as a tuple M = ⟨S, A, R, T , γ⟩, where S is the set of states, A the set of actions, R is the reward function, T represents transition dynamics, and γ is a discount factor. The solution represented by an optimal policy is defined as the one maximizing the expected discounted cumulative reward of the trajectory:EQUATIONHere, γ is a discount factor, transition probability p(s ′ |s, a) is fixed but unknown, and the policy is learned to optimize this objective from experience. Now we consider MDP setting, where in addition to transition probabilities, the task reward is unknown. Our objective is to reconstruct this reward function from observational data. As proposed in [10] , we represent the reward as a linear combination of features which can consists of any nonlinear function of the state.EQUATIONwhere s is the state, ϕ ∈ R d is the feature function, mapping the state s to a d-dimensional feature vector, and θ ∈ R d is the weight vector of features. The reward construction problem then translates to finding suitable features and their corresponding weights. 
proposed method: 
 In this section, we present our method for learning a reward model, which consists of two aspects: finding the feature set and its weight vector. As ground truth rewards are not available, we use unsupervised feature selection by generating pseudo labels from expert data. Then, we adopt the max-ent IRL framework for learning feature weights. The overview of the algorithm is provided in Appendix C.According to the maximum entropy formulation for stochastic policies [19] , the probability of a trajectory is proportional to the exponential of the reward of that trajectory. That is,EQUATIONwhere the partition function Z(θ) normalizes the reward function over all possible trajectories. After taking logarithm of both sides of (3), we getEQUATIONand noting that the partition term Z is constant:EQUATIONDue to the linearity from Eq. 2, the cumulative reward of a trajectory can then be defined asEQUATIONwhere ϕ(τ ) = s i ∈τ ϕ(s i ). Thus, the log probability of trajectories is proportional to a linear combination of features that compose the reward.EQUATIONBased on this proportionality, to determine whether a specific feature contributes to the reward, we evaluate the correlation between feature's expectation and log probability of trajectories. Since ground truth rewards are not available, we use log probabilities of trajectories as pseudo-labels.We generate pseudo labels by computing log probabilities of trajectories empirically from training dataEQUATIONWe simplify the computation of conditional probabilities by decomposing them into the marginal and joint probabilities of consecutive states. After applying the logarithm to both sides and simplifying the terms, we getEQUATIONThe derivation of the Eq. 14 is included in the Appendix D.1. Using the entire training dataset, we perform multivariate kernel density estimation (KDE), to obtain marginal P (s) and the joint density of sequential states P (s t , s t+1 ). KDE provides a non-parametric way to estimate the probability density function of the states and requires substantial less training data, unlike neural networks.To perform feature selection using Eq.7, a finite set of candidate state features is required, as evaluating all possible features may be computationally infeasible. According to the moment matching method [41] , the probability of the state s can be approximated by their higher moments. Using moments up to the third order, we can approximate state probability aslog P (s) ≈ Φ 1 (s) + Φ 2 (s) + Φ 3 (s)where Φ 1 (s) represents the first-order moments (mean), Φ 2 (s) represents the second-order moments (covariance), Φ 3 (s) represents the third-order moments (skewness), s is a state vector. As proposed features are evaluated on their predictive power of the log probabilities of trajectories, which in turn consist of state trajectories, we include all covariance terms in the candidate set of features Φ(s) ∈ {Φ 1 (s), Φ 2 (s), Φ 3 (s)}. However, due to the curse of dimensionality and the finite set of demonstrations, using all features as reward components is impractical. Feature matching with IRL may fail due to noise and spurious correlations, leading to false correlations that do not exist in the original data distribution. Therefore, we present an efficient method to select only the relevant features.After computing the probability and feature expectations of trajectories, we obtain the following data.X =     ϕ 1 (τ 0 ) ϕ 2 (τ 0 ) • • • ϕ K (τ 0 ) ϕ 1 (τ 1 ) ϕ 2 (τ 1 ) • • • ϕ K (τ 1 ) . . . . . . . . . . . . ϕ 1 (τ N ) ϕ 2 (τ N ) • • • ϕ K (τ N )     Y =     log P (τ 0 ) log P (τ 1 ) . . . log P (τ N )    One way to perform feature selection is to utilize a univariate feature selection method based on statistical tests. First, for each feature separately, we apply F-statistics between trajectory feature expectations and trajectory probabilities, i.e., between each column of X and the vector Y . This gives the feature importance metric, which we can use for ranking the features. Then, the features with higher F-statistics are selected and included in the feature extractor ϕ(•). Univariate feature selection is more suitable when the number of samples is much smaller (K ≫ N ) than the size of the candidate set, which is often the case for high-dimensional states. Therefore, directly using linear regression for X and Y results in overfitting. Recursive feature elimination methods may be used here, but they are computationally demanding. Although the univariate feature selection method does not account for interactions between features, these interactions will be considered while learning the reward function through inverse reinforcement learning. Additionally, to reduce overfitting, we employ crossvalidation and noise addition methods. Specifically, we add multiple non-expert trajectories by labeling them with uniformly sampled log probability values from the bottom 10th percentile of expert log probabilities. The time complexity of the feature selection algorithm is O(N ) with respect to the size of the candidate set Φ. This approach is efficient, as it does not involve iterative search of features and policy learning, which would be the case with manual feature construction. The next step to fully recover the reward function is to find the weights of each of the features, which is covered in the next part.We employ the maximum entropy IRL [19] to learn the weights of the selected features. To maximize the log probability of the observed data, we formulate the optimization problem as follows:EQUATIONTo find the optimal weights θ, we take the derivative of log-likelihood that is given by:EQUATIONHere, µ e represents the feature expectation of the training data, m is the number of trajectories from current policy. The weights of the reward can then be updated by using the gradient descent as follows:EQUATION)where α is the learning rate. The full derivations are included in Appendix D.2. The learned reward function has an explicit form allowing us to apply any value iteration algorithm to obtain a policy. We utilize the Soft Actor-Critic (SAC) [42] to learn the expert policy by optimizing Eq. 1. 
experiments: 
 In this section, we consider different continuous control tasks and learn their reward functions from expert data. Then, using the learned rewards, we learn the policy and compare it against the expert performance. We will investigate the following research questions:1. Is the proposed methodology effective in reconstructing ground truth rewards solving various robotic tasks?2. Are the learned rewards transparent and amendable? We conduct experiments on different environments with continuous state space, particularly Mu-JoCo environments such as Ant, HalfCheetah, Walker2d, and Hopper from the Gymnasium [43] benchmark (see Fig. 2 ). The descriptions of the tasks and ground truth reward formulations are provided in Appendix A.To collect the expert or training data, we train RL algorithms for the above environments. Specifically, we used SAC for all the environments with a continuous action space. We train RL algorithms to maximize the cumulative ground truth reward function of the environment. After reaching established benchmark results as presented in Stable Baselines3 [44] for these tasks, we refer to the RL policy as an expert policy and deploy it for data collection. For each of the environments, we collected N trajectories with random starting states, and the simulated trajectories were saved in dataset D. Thus, in our IRL method, as shown in Algorithm 1, only dataset D is used, and we assume that the expert policy or ground truth reward function is unknown. We trained the SAC algorithm for a very limited time-40,000 simulation steps-to generate sensible but suboptimal(non-expert) trajectories for feature selection phase. More implementation details are provided in Appendix B.In this work, we propose a feature selection mechanism aimed at recovering the reward function that corresponds to the expert's behavior, given the dataset. Therefore, we compare our method against various baseline feature selection strategies: manual features, randomly selected features, direct use of all states as features (referred to as linear features). For manual feature construction, we evaluated several hand-crafted features and selected the best ones for final experiment. Similarly, the random projection method evaluates a random subset of features from the candidate set. Given that comparing reward functions directly across baselines does not yield a meaningful performance metric in terms of completing a task, we focus on comparing the policies derived from these reward functions.To this end, we conduct two comparative analyses. First, we execute the derived policies in multiple testing environments configured by ground truth reward functions and observe cumulative rewards. This evaluation is valid because our training data comes from an expert trained with ground truth reward function as well. Second, we compare the state distributions of the training data from expert and testing data from the extracted policies. For the divergence measure of multivariate distributions, we exploited the 2D Wasserstein distance metric [45] . Also, we compared the performance of the method against SOTA adversarial inverse reinforcement learning methods such as AIRL [27] and GAIL [35] . The cumulative episodic rewards of policies trained using different rewards using baseline methods are shown in Table 2 .The results show the average for 10 trajectories with three random seeds. The per- formance of the proposed reward model outperforms the baselines for all of the tasks.Next, we compare the performance of our method to AIRL and GAIL algorithms from the imitation learning library [46] . Table 3 presents the mean cumulative rewards for the tasks studied. We can see that our method achieves higher cumulative rewards in all domains, with significant improvement on the Walker2d, Hopper, and Ant tasks. While adversarial IRL methods achieve state-of-the-art (SOTA) performance, they require a large amount of training data to perform well. Table 1 shows the number of parameters in the reward model. In adversarial IRL methods, a two-layer MLP with 32 nodes is used, while proposed method uses linear reward model with significantly less parameters. This enables to simplify learning process of the reward and therefore, proposed method performs better than adversarial methods with less data. Expert Fast Slow vx 0.12 ± 0.1 0.17 ± 0.18 0.09 ± 0.08 Table 4 : Velocity of the HalfCheetah along the x-axis changes accordingly after modifying the weights of the reward component s 8 .The learning curves of the IRL for Walker2d are shown for different reward models in Fig. 3a . Similarly, the 2D Wasserstein distance metric [45] between expert and learner state distributions during training is shown in Fig. 3b . To compute the divergence metric, we use 40 trajectories of expert and learner. The divergence metric between expert and learner policies clearly demonstrates that for the proposed reward, not only does the policy achieve high cumulative rewards, but it also better matches the expert state distribution compared to other baseline methods. Next, Fig. 4a shows the correlation plot between ground truth rewards and recovered rewards for 150 trajectories with random initial conditions. The Pearson correlation score is reported in the figure for the HalfCheetah environment. We observe a strong correlation between recovered and ground truth rewards, with a Pearson correlation coefficient of 0.78. This finding is noteworthy, as no ground truth reward signal was provided during reward learning; only expert demonstrations, specifically state trajectories, were used to infer rewards. This suggests that the obtained rewards not only result in behaviors similar to those of the expert but also achieve this by constructing a reward signal that closely resembles the ground truth reward signal. In addition, Fig. 4b shows the full RL training loop with the recovered reward, suggesting the compatibility of the recovered rewards with the RL framework, meaning that RL can be trained from scratch using this reward function. The transparency of the reward can be seen from the reward representation which consists of monomials. This enables the analysis of which features of the expert contributes to reward and interaction between state features. We observe that the reward terms focus on state s 8 , which represents the velocity along the x-axis, directly relating to the agent's original reward. Additionally, there is a high negative gain for s 9 , the angular velocity along the y-axis, which is sensible since the HalfCheetah should primarily move along the x-axis. Negative gains for the angular velocity of the rotors, specifically s 2 and s 16 , are also evident, aligning with the control penalty in the ground truth reward.One important benefit of this compact and explicit representation of the reward is the possibility of fine-tuning the reward function to generate different behaviors. The Table 4 below shows the mean and standard deviation of the HalfCheetah's velocity along the x-axis after multiplying coefficient of one of the terms, s 8 , in the reward function by 2 and 0.5, respectively. As a result of amending reward components, faster and slower motion was achieved, which can be expected as state s 8 accounts for the velocity along the x-coordinate of the front tip.One limitation of this paper is its applicability to other high-dimensional inputs, such as images. We anticipate that pre-learning features, possibly through the use of autoencoders, and combining with proposed method could be a viable approach to overcoming this limitation, and we plan to explore this direction in future work. Another limitation is our experimental design. We have not conducted experiments with real robots and human demonstrations. However, we validated our method using complex robotic environments. In future work, we plan to apply the method to real-world tasks and learn rewards from demonstrations. 
conclusion: 
 This paper presented a method for constructing a transparent reward model with an explicitly represented form using non-linear state features. Reward components were selected using an unsupervised learning method from observational data by incorporating the maximum entropy formulation for the behavior. As shown above, selecting the features that allow predicting trajectory probabilities results in a higher match with ground truth rewards and enables the retrieval of policies more similar to those of an expert. Learning rewards from data alleviates the challenge of manual reward specification for complex tasks and allows for training with exploration outside of the training data using the reinforcement learning framework. We have demonstrated that across different environments, recovered rewards outperform other baseline methods. Table 5 : State-action dimensionality of the tasks An additional challenge in reward learning is that the features composing the ground-truth rewards are not directly available. For example, the direct components of the ground-truth reward, particularly the x-coordinate and torque values, are hidden from the states and, consequently, from the reward model. Instead, we infer the reward from the available indirect features like velocity or joint angles. 
 
 Appendix: The source code was implemented in Python 3.8, and the source code will be publicly available upon acceptance. We utilized the Stable-baselines3 library [44] for the training of RL algorithms. As a policy network, we used a multi-layer perceptron (MLP) with two hidden layers. Furthermore, we conducted a thorough hyperparameter optimization for both the RL and inverse reinforcement learning (IRL) parameters. A table with the hyperparameters used during the training is included in the supplementary materials. The total number of training iterations for IRL is 50, and the training dataset consists of 150 trajectories. To facilitate the training, we parallelized data collection and environment simulation using multi-processing techniques. During the feature selection process, we normalize each state feature to have a mean of zero and a standard deviation of one. We apply the same normalization for reward computation. However, for policy learning, we use the raw state inputs. We employed the Optuna [47] library with grid search to tune hyperparameters, and Adam optimizer [48] was used to learn reward model. After applying the logarithm to both sides and simplifying the terms, we get log P (τ ) = log P (s 1 ) + n-1 t=1(log P (s t , s t+1 ) -log P (s t ))Here, conditional state probability P (s t+1 |s t ) can be written as P (s t+1 |s t ) = P (s t , s t+1 )/P (s t ). Then: P (τTo find the optimal weights θ, we take the derivative of log-likelihood that is given by: EQUATION≈ µ e -m i=1 s i ∈τ ϕ(s i )HalfCheetah  


Decision: **Summary**

The paper introduces a method for learning transparent reward functions from human demonstrations by focusing on selecting important state features. Using a maximum entropy framework, the approach involves estimating trajectory density and identifying the features that best model this density. This feature selection method for IRL aims to extract state features that are highly correlated with expert behavior, enhancing the transparency of the learned reward functions.

**Summary of Reviews**

*Strengths*:
- Reviewers found the idea of automatically selecting state features for reward models compelling, novel, and principled.
- The experiments were well-designed to demonstrate the benefits of the proposed approach.

*Weaknesses*:
- Lack of clarity. All reviewers criticized the paper for its lack of clarity, missing crucial details necessary for understanding the method, the baselines, and experimental details.
- Missing technical and experimental details. All reviewers had lingering questions about the method and experiments. The authors should resolve any remaining confusions in the rebuttal.
- Insufficient experiments. Some reviewers complained about the narrow evaluation on just the Hopper, Ant, Walker and HalfCheetah environments. Others asked for more baseline comparisons to AIRL and GCL. The lack of real-world robotics experiments is concerning.
- Missing discussion of limitations.

**Outcome and Post-rebuttal**

Two reviewers substantially improved their scores and indicated their concerns being mostly if not fully addressed. I thank the authors for such a substantial and compelling rebuttal effort! Reviewers found the manuscript significantly improved, and appreciated the reframing as transparency and not interpretability. I recommend Accept for poster and I suggest the authors follow the remaining reviewer suggestions to improve their manuscript for the camera ready. In particular: more details for reproducibility, more details about the handcrafted features, a comparison with FIRL, and a code release.

Reviews: review: title: title: Interesting research on feature selection for reward learning but experimental results are not very impressive.
paper_summary: paper_summary: For learning a reward from human demonstrations, the authors propose a new method to select important state features. Based on max entropy formulation, the authors use trajectory density as labels and find best features to model the trajectory density.
main_review: main_review: **Strengths**

- S1: Interesting research direction. The idea of automatically selecting state features to develop effective and interpretable reward models is compelling. To achieve this, the authors proposed a new solution based on trajectory density estimation from expert demonstrations.

- S2: Well-designed experiments. The experiments are thoughtfully designed to demonstrate the benefits of the proposed method. In addition to policy learning results, the authors provide insightful analysis on interpretability and state distribution matching.

**Weaknesses**

- W1: The authors used simple simulated robots like Hopper, Ant, Walker, and HalfCheetah in their experiments. This limits the impact of the experimental results. It would be beneficial if the authors could show results using robots with more joints (such as Humanoid), which would involve a higher-dimensional state space.

- W2: Several components are hard to understand: (1) The explanation in lines 98-100 is a bit unclear. This section is crucial for explaining the motivation and intuition behind using log probabilities as pseudo-labels. It would be helpful if the authors could elaborate more on this point. (2) The process of creating candidate features is unclear. The authors introduced moment matching, but it is not clear what $\phi$ represents exactly. More details on this aspect would be beneficial.
questions: questions: **Random seeds**

- The authors mentioned, "The results show the average for 10 trajectories with random seeds." How many random seeds were used?

**Writing**

- line 112: change "states ss" to "state $s$"

- Eq3,4,5,7: change $P(\tau_i)|\theta)$ to $P(\tau_i|\theta)$

- Algorithm 1, line6: change $(D,\phi$ to $(D,\phi)$

- Eq 1: $t$ notation is missed inside of expectation.
limitations: limitations: 3
review_summary: review_summary: The paper explores an interesting direction, but some details are not clear. Additionally, the experimental results are somewhat limited. Therefore, I recommend a 'weak reject.'
title: title: Review
paper_summary: paper_summary: This paper seeks to learn interpretable reward functions from demonstrations
main_review: main_review: This paper studies the important problem of learning intepretable reward models. The motivation is well done and the idea of using log probabilites of trajectories to find an appropriate set of candidate features is novel as far as I can tell.

Is there a reason to estimate the marginals and joint probability distributions as opposed to estimating the conditional distributions?

Noise addition method is unclear. How do you use non-expert trajs? More details would be good.

The paper states that the method does not include many combinations of features. It's unclear what this means and whether it is a good or bad thing. It seems that you would want to be able to learn non-linear features that are combinations of state features.

Eqns 11-14 seem to be identical to the Max Ent original paper so you could cut space there by not including the derivation.

It would be nice to have more details on the manual feature construction baseline as there are many ways to do this. How were the features determined? What kind of trial and error process was used? 

Furthermore, it's unclear what kinds of features your method uses. Are they just raw state dimensions? Is there any preprocessing? How would this work with visual inputs where you wouldn't want to have each pixel be a feature? It seems that the reward is just a function of raw state, but the reward shown on line 228 is a non-linear function of state and it's unclear to me how this type of reward is obtained. Have the authors considered feature learning approaches that return a linear reward function such as 
- Brown et al. "Safe imitation learning via fast bayesian reward inference from preferences." ICML, 2020.

Figure 2 is nice in showing the behavioral similarity between the expert and the learned policy when using the proposed approach.

Why only compare against AIRL on two domains? It would be good to include the AIRL comparison in the Table 1 across all domains. Also other baseline approaches such as "Feature construction for inverse reinforcement learning" would be good to compare against if possible.

The idea of interpretable rewards is very appealing, but looking at the half-cheetah reward on line 228 it is hard to tell how interpretable it is (it doesn't seem very interpretable to my untrained eye). It would be nice to actually study the interpretability and to compare with other interpretability approaches for reward learning.

Overall the results in Table 1 and 2 show that the proposed approach works better than other approaches, but some of the baselines are unclear such as the manual and random projection. Comparing with AIRL and maybe Guided cost learning would be good as well as comparing against other feature selection/interpretable reward learning approaches.





Several typos and formatting issues throughout.
- log P(\tau_i) | \theta) ->  log P(\tau_i | \theta)
- line 112 "states ss" 
- Reference [27] has the wrong authors for the AIRL paper.
- line 236 has a spurious newline
questions: questions: See questions above.
limitations: limitations: 1
review_summary: review_summary: The paper studies an important area and has some promising results but there are insufficient details to fully evaluate and the claim of interpretability requires more validation.
title: title: Not ready for publication
paper_summary: paper_summary: This paper proposes a feature selection method for IRL which extracts state features and values highly correlated with the expert behavior through an explicit feature extraction and density estimation step.
main_review: main_review: ## Strengths

**Method:** The method seems principled and incorporates traditional machine learning techniques for deep IRL

**Results:** Results are not bad, demonstrating the method is better than ablations and (somewhat unconvincingly with no standard deviations) that the method possibly outperforms AIRL.

## Weaknesses

**Writing:** I have quite a few complaints about some of the writing and clarity of the paper:

- The first paragraph of the intro doesn’t hint at or make clear what problem the paper is trying to solve.
- Combined with the above point, there are no method or teaser figures that clearly demonstrate the problem setting (inputs/outputs) and method at a high level. This makes it hard to understand what the paper will be about until the end of the intro.
- Interpretability isn’t talked about much throughout the paper yet it’s highlighted in the title and briefly mentioned in the intro. More justification should be given for how this method is more interpretable than other methods that learn a linear function of the state features as a reward.
- Notation: the notation for how to write conditional probabilities seems highly nonstandard: $P(\tau_i)|\theta)$ is written many times but to me it just looks like there’s an extra parenthesis after $\tau_i$.
- Section 4.2: why kernel density estimation in particiular? no intuition is included of using this versus any other method of estimating probabilities (e.g., train a neural network)
- Section 4.3: Why the 3rd order specifically? the experiments don’t justify this choice over 4th order or just 2nd order?

**Experiments:** 

- How many seeds were used? This should be specified
- The AIRL baseline is from 2017, even though it’s a valid baseline. Perhaps a few more recent baselines would make sense to compare to.
- There's no standard deviations for the AIRL baseline either, so these results are unconvincing. Furthermore, it's only presented on half the tasks.
- Figure 2: what’s an epoch on the x axis? I’d recommend remaking these WandB plots using their API to download the stats and plotting them with proper labels and bigger font sizes.
- There’s some missing analysis in the experiment results, e.g. L 241: why is the proposed method better than AIRL here?
- How are hyperparameters tuned for the proposed method and ablations/baselines?
- Not much is done with respect to the interpretability claim. One could perform experiments demonstrating that the form of interpretable reward functions is actually amenable to editing/preferred by humans/etc, e.g., Section 5.7 of this [paper](https://clvrai.github.io/leaps/).

Limitations: The authors should include some discussion of the limitations in the paper.




POST REBUTTAL:

A lot of my concerns have been addressed, and after looking at the revised paper, I am updating my score to a weak accept. I think the paper has merit.
questions: questions: See above questions in weaknesses.
limitations: limitations: 1
review_summary: review_summary: Not ready due to some substantial writing issues, a lack of new baselines, and insufficient experiments.