Automated Feature Selection for Inverse Reinforcement Learning

Abstract: Inverse reinforcement learning (IRL) is an imitation learning approach to learning reward functions from expert demonstrations. Its use avoids the difficult and tedious procedure of manual reward specification while retaining the generalization power of reinforcement learning. In IRL, the reward is usually represented as a linear combination of features. In continuous state spaces, the state variables alone are not sufficiently rich to be used as features, but which features are good is not known in general. To address this issue, we propose a method that employs polynomial basis functions to form a candidate set of features, which are shown to allow the matching of statistical moments of state distributions. Feature selection is then performed for the candidates by leveraging the correlation between trajectory probabilities and feature expectations. We demonstrate the approach's effectiveness by recovering reward functions that capture expert policies across non-linear control tasks of increasing complexity. Code, data, and videos are available at https://sites.google.com/view/feature4irl.

• The paper presents a method that uses polynomial basis functions to construct candidate features by matching key statistical moments of state distributions.
• It employs a correlation-based feature selection process that effectively reduces noise and complexity by ranking candidate features.
• Experiments in Gymnasium environments demonstrate that the approach replicates expert behavior using fewer features compared to baseline methods.

Automated Feature Selection for Inverse Reinforcement Learning

Introduction

Inverse Reinforcement Learning (IRL) provides a powerful framework for learning reward functions from expert demonstrations, eliminating the necessity for explicit reward specification while retaining the generalization advantages inherent to reinforcement learning (RL). The reward within IRL is traditionally represented as a linear combination of features, which challenges researchers to identify suitable, rich features, especially in continuous state spaces where mere state variables are inadequate. The paper presents a method utilizing polynomial basis functions to construct a set of candidate features, which statistically match the moments of state distributions and leverage correlations between trajectory probabilities and feature expectations for optimal feature selection.

Figure 1: A central open challenge in inverse reinforcement learning is the choice of suitable features to represent the reward. We propose a method that constructs a candidate feature set and then selects a subset that best describes expected rewards.

Methodology

Polynomial Feature Selection

The method initiates by generating candidate features as quadratic polynomial functions of the state. This choice facilitates the matching of the statistical moments, specifically the mean and variance, between the expert demonstrations and the policy obtained, aligning the Gaussian approximations of state distributions while maintaining a manageable feature dimensionality.

Feature Selection Mechanism

An efficient feature selection process is implemented that employs correlation-based techniques to identify features with the highest relevance by leveraging the relationship between trajectory probabilities and feature expectations. This process considerably reduces reward complexity and effectively addresses noise or spurious correlations. The algorithm ranks features using statistical tests and selects those deemed most promising, optimizing for compactness and interpretability.

Reward and Policy Retrieval

Maximum entropy IRL is adapted to discern the weights of selected features, optimizing the feature weights by aligning the feature expectations of expert trajectories with the policy-derived ones. This involves gradient descent optimization over the computed feature expectations using RL (PPO and SAC algorithms) to extract an optimal policy that maximizes the expected cumulative reward based on the retrieved reward function.

Experiments

The experimental setup encompassed three Gymnasium environments—Pendulum-v1, CartPole-v1, and Acrobot-v1—allowing the rigorous evaluation of the proposed method by comparing the retrieved policies against known benchmarks.

Figure 2: Pendulum-v1.\ dim(Phi) = 9

Key findings illustrated how the proposed method achieved benchmark results while employing fewer features compared to baselines like hand-picked features, random selection, and the inclusion of all candidate features. The superiority was evident in the effective replication of expert behavior across varying complexity environments, confirming the efficacy of polynomial basis functions in overcoming non-linear control challenges.

Figure 3: Mean cumulative rewards for policies trained using various feature sets, calculated across 10 different initial conditions. A) Pendulum, B) Acrobot, C) CartPole.

Discussion

The research advances automated IRL feature selection efficiency and accuracy, emphasizing the prospect of incorporating diverse basis functions such as radial basis functions (RBFs) and Fourier series. This generalization potential positions the method for broader applicability, particularly in preference learning scenarios among tasks with consistent environmental setups and varying expert influences.

Figure 4: 2D Wasserstein distance between training and testing data for the Pendulum and Acrobot environments.

Conclusion

This paper introduces refined algorithms that employ polynomial basis functions for constructing reward functions in inverse reinforcement learning, demonstrating substantial effectiveness in diverse environments. By facilitating the streamlined selection of pertinent features, the proposed method enhances both process efficiency and model interpretability, setting the stage for further refinement with alternative basis functions to boost precision and broaden applicability. Future endeavors are expected to build on these foundations, exploring more sophisticated strategies to enhance the adaptability and efficacy of automated feature selection in imitation learning contexts.