[Paper Review] k-Means Maximum Entropy Exploration

Alexander Nedergaard, Matthew Cook|arXiv (Cornell University)|May 31, 2022

Neural dynamics and brain function4 citations

TL;DR

This paper proposes k-Means Maximum Entropy Exploration (KME), a computationally efficient artificial curiosity algorithm that uses a k-means-based lower bound on state visitation entropy to enable exploration in high-dimensional, continuous reinforcement learning environments. The method leverages a theoretical result on non-parametric density estimation via k-means to approximate entropy, showing strong performance on sparse-reward benchmarks, especially where standard RL fails to find rewards.

ABSTRACT

Exploration in high-dimensional, continuous spaces with sparse rewards is an open problem in reinforcement learning. Artificial curiosity algorithms address this by creating rewards that lead to exploration. Given a reinforcement learning algorithm capable of maximizing rewards, the problem reduces to finding an optimization objective consistent with exploration. Maximum entropy exploration uses the entropy of the state visitation distribution as such an objective. However, efficiently estimating the entropy of the state visitation distribution is challenging in high-dimensional, continuous spaces. We introduce an artificial curiosity algorithm based on lower bounding an approximation to the entropy of the state visitation distribution. The bound relies on a result we prove for non-parametric density estimation in arbitrary dimensions using k-means. We show that our approach is both computationally efficient and competitive on benchmarks for exploration in high-dimensional, continuous spaces, especially on tasks where reinforcement learning algorithms are unable to find rewards.

Motivation & Objective

• To address the challenge of efficient exploration in high-dimensional, continuous reinforcement learning environments with sparse extrinsic rewards.
• To provide a principled, mathematically grounded approach to intrinsic curiosity by maximizing the entropy of the state visitation distribution.
• To overcome the computational infeasibility of estimating state visitation entropy in high dimensions using standard non-parametric methods.
• To develop a practical algorithm that approximates maximum entropy exploration using k-means clustering with provable theoretical underpinnings.

Proposed method

• Proposes a novel lower bound on the entropy of the state visitation distribution using k-means clustering in arbitrary dimensions.
• Derives a theoretical result linking k-means cluster diameters and density estimation, proving that balanced Voronoi cells yield consistent entropy estimates in the limit of infinite clusters.
• Constructs an intrinsic reward signal based on the inverse of cluster diameters and the number of clusters, approximating the entropy lower bound.
• Uses a practical variant of k-means that assumes balanced Voronoi diagrams and replaces the logarithmic term in entropy with a square root for stability and efficiency.
• Employs the entropy lower bound as an intrinsic reward in a reinforcement learning loop to encourage uniform state visitation.
• Validates the method on continuous control benchmarks, including environments where standard RL agents fail to discover sparse rewards.

Experimental results

Research questions

• RQ1 Can k-means clustering provide a scalable and theoretically grounded approximation to the entropy of the state visitation distribution in high-dimensional continuous spaces?
• RQ2 Does a lower bound on entropy derived from k-means cluster geometry lead to effective exploration in sparse-reward reinforcement learning tasks?
• RQ3 How does the proposed k-Means Maximum Entropy Exploration (KME) method compare to existing intrinsic curiosity and count-based exploration methods in high-dimensional environments?
• RQ4 What are the practical limitations of the method in very high-dimensional tasks, such as Humanoid?
• RQ5 Can the method be improved by combining it with representation learning or better cluster initialization strategies?

Key findings

• The proposed k-Means Maximum Entropy Exploration (KME) method achieves competitive performance on benchmarks for exploration in high-dimensional, continuous control tasks, particularly in scenarios where standard RL agents fail to find sparse rewards.
• The method demonstrates computational efficiency due to the use of k-means clustering and a stable approximation of the entropy lower bound.
• Empirical results show that the entropy estimate preserves the correct entropy ordering across different distributions (e.g., uniform, Gaussian mixtures, Gaussians) in 2D, 4D, and 64D, though convergence is slow in high dimensions.
• The algorithm struggles to find rewards in the Humanoid environment, indicating limitations in very high-dimensional state spaces.
• Theoretical analysis supports the validity of the entropy approximation under idealized assumptions, and empirical results suggest the practical approximation is sufficient for effective exploration.
• Slow convergence in high-dimensional random walks is attributed to poor initialization of cluster centers at zero, which can be mitigated with better initialization strategies.