SEMDICE: Off-policy State Entropy Maximization via Stationary Distribution Correction Estimation

Reinforcement learning has achieved notable success in complex tasks like games and robotics, yet relies heavily on reward signals. Unsupervised pretraining methods, such as exploration via intrinsic rewards, aim to address data inefficiency and generalization. State entropy maximization (SEM) has emerged as a promising exploration strategy, encouraging agents to visit diverse states. However, existing methods rely on policy-gradient approaches or particle estimators, which are either sample-inefficient or biased. Recent advances like DICE algorithms have improved off-policy distribution estimation but have not been applied to maximizing state entropy directly. The challenge remains to develop an unbiased, sample-efficient off-policy method capable of estimating and optimizing the state distribution entropy, especially from arbitrary datasets.

The core problem is estimating and maximizing the entropy of the target policy’s stationary state distribution using off-policy data. Existing methods often estimate the entropy of stored experiences, leading to bias and inefficiency. Importance sampling can correct off-policy bias but suffers from high variance, especially over long horizons. Consequently, current algorithms either lack theoretical guarantees or are too sample-inefficient for practical use. Developing an unbiased, stable, and efficient off-policy method for state entropy maximization is crucial for scalable unsupervised RL pretraining, enabling agents to explore more effectively and adapt rapidly to downstream tasks.

This paper introduces SEMDICE, a novel approach that formulates the problem as a convex optimization over stationary distributions, directly maximizing state entropy with regularization. Key innovations include: 1) leveraging f-divergence for distribution correction, 2) transforming the constrained problem into an unconstrained dual form via Lagrangian methods, 3) approximating expectations with Monte Carlo sampling, and 4) parameterizing dual variables with neural networks for model-free training. Unlike prior particle-based or policy-gradient methods, SEMDICE guarantees unbiased estimation and convergence to the optimal maximum entropy distribution. Its ability to learn from arbitrary off-policy datasets without bias or high variance marks a significant advancement in off-policy RL.

• �� Formulate the state entropy maximization as a convex programming problem over stationary distributions, incorporating Bellman flow constraints and f-divergence regularization. • Use Lagrangian duality to convert constrained optimization into an unconstrained convex problem, optimizing dual variables (ν, μ). • Approximate expectations via Monte Carlo sampling, parameterize ν and μ with neural networks for flexibility. • Derive the distribution correction ratio w* from the dual solution, enabling unbiased estimation of the target state distribution. • Use a single convex minimization step to find the optimal distribution, ensuring convergence. • Extract the policy by normalizing the optimal distribution ratios, employing an i-projection method. • During training, alternate updates of ν, μ networks, and policy parameters, ensuring stability and convergence. • Theoretically, the method guarantees convergence to the maximum entropy stationary distribution, validated through experiments.

Experiments involve randomly generated finite MDPs with 20 states and 4 actions, comparing SEMDICE to baselines like count-based exploration and particle estimators. Metrics include state entropy and sample efficiency, evaluated over 100 runs with different seeds. Additionally, URL benchmark tasks assess downstream adaptation speed and performance after pretraining. Hyperparameters like regularization coefficient α and divergence choice are tuned via ablation studies. Results show SEMDICE consistently converges to the optimal state entropy, outperforming biased and on-policy methods, with faster adaptation and higher exploration coverage. The experiments validate theoretical guarantees and demonstrate practical robustness across environments.

SEMDICE achieves over 15% higher state entropy than baseline methods in finite MDPs, converging reliably across multiple seeds. In URL benchmarks, pretraining with SEMDICE accelerates downstream task adaptation by approximately 30%, especially in sparse reward settings. Ablation studies reveal that tuning the regularization parameter α and divergence type significantly influences exploration quality, balancing bias and variance. The results confirm that directly optimizing in the stationary distribution space yields superior exploration and transfer performance, validating the theoretical advantages of the proposed convex approach.

This method is suitable for autonomous exploration in robotics, where agents need to learn diverse behaviors without explicit rewards. It can be integrated into pretraining pipelines for complex control tasks, enabling rapid adaptation to new environments. The approach also benefits applications like autonomous driving, where exploration and transfer learning are critical. Its ability to learn from arbitrary off-policy datasets makes it highly versatile for real-world scenarios with limited online interaction.

The current implementation relies on neural network parameterization, which may require extensive tuning and computational resources. Its theoretical guarantees are primarily established for finite or low-dimensional spaces; scalability to high-dimensional continuous environments remains to be demonstrated. The sampling complexity and training stability could be challenged in very large or complex environments, limiting real-time deployment. Future work should focus on improving scalability, reducing computational costs, and extending theoretical analysis to broader settings.