Skip to content

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

🕒 Published (v1): 2025-12-10 19:50 UTC · Source: Arxiv · Venue: ICLR 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

SEMDICE is a principled off-policy algorithm for state entropy maximization (SEM) in unsupervised RL pre-training. It reformulates SEM as a single convex minimization problem over the space of stationary distributions using DICE-style stationary distribution correction, enabling unbiased learning from arbitrary off-policy data. It outperforms prior SEM methods on convergence and downstream task adaptation efficiency.

Problem

Existing SEM pre-training methods are either on-policy (sample-inefficient) or off-policy but biased: they estimate the entropy of the replay buffer's state distribution rather than the target policy's stationary state distribution. Importance-sampling corrections suffer from high variance due to the curse of horizon. No prior method offers both unbiasedness and sample efficiency for SEM.

Method

SEMDICE formulates the regularized SEM objective directly over state-action stationary distributions d(s,a) with an f-divergence regularizer constraining deviation from the off-policy dataset distribution. Via Lagrangian duality and Fenchel conjugation, the constrained concave program reduces to a single unconstrained convex minimization over dual variables ν(s) (Bellman multipliers) and µ(s) (marginalization multipliers). The objective is made numerically stable by replacing the exp sum with a log-sum-exp (softmax-normalized) form. A Monte Carlo approximation over the replay buffer's empirical state distribution handles continuous state spaces. The resulting stationary distribution correction ratios w*(s,a) = d*(s,a)/dD(s,a) are parameterized via MLPs; the explicit policy is extracted via i-projection (KL minimization against the reweighted behavioral policy).

Key Contributions

  • First principled, unbiased off-policy SEM algorithm: reduces SEM to a single convex minimization over stationary distributions.
  • Derivation via Lagrangian duality + Fenchel conjugate that eliminates the inner max and yields a tractable objective from arbitrary off-policy data.
  • Numerically stable objective (Theorem 3.3) replacing exp sum with log-sum-exp without changing the optimal solution up to a constant shift.
  • Empirical demonstration of convergence to the optimal SEM policy in tabular MDPs where all prior baselines fail.
  • State-of-the-art SEM-based pre-training on URLB across 12 locomotion and manipulation tasks.

Results

  • Tabular MDPs (20 states, 4 actions, 100 seeds): Only SEMDICE converges to the optimal SEM policy; baselines (CB-SA, CB-S, PB-S) do not. On-policy PG requires ×100 more samples.
  • MountainCar visualization: SEMDICE achieves progressive and stable state coverage convergence; SEM baselines (ProtoRL, APT, MEPOL) exhibit oscillatory/non-stationary behavior.
  • URLB pre-training entropy (Walker, Quadruped, Jaco Arm, 2M steps): SEMDICE achieves the highest particle-based entropy across all three domains; reaches maximum entropy in Walker and Jaco with as few as 100K interactions vs. ~2M for competing methods.
  • URLB fine-tuning (100K downstream steps, 120 runs total): SEMDICE achieves the best IQM, Median, Mean, and Optimality Gap among all data-based and knowledge-based methods; underperforms competence-based CIC overall.

Limitations

  • Underperforms CIC (competence-based method using mutual information + representation learning); authors attribute this to lack of representation learning components.
  • The single-sample estimate of the Bellman term introduces bias (upper-bounds the true objective); equality holds only for deterministic MDPs.
  • No formal convergence guarantee for the continuous/function-approximation case.
  • Not evaluated on pixel-based or high-dimensional observation domains; representation learning extension is left to future work.
  • The f-divergence regularization hyperparameter α requires tuning.

Relevance to Agentic AI / LLM Agents

SEMDICE addresses unsupervised pre-training of agents that must generalize to unknown downstream tasks — a core challenge in building autonomous agents that operate without hand-specified reward functions. The principle of optimizing directly over stationary distributions using off-policy data is analogous to how LLM-based agents pre-train broadly before task specialization, and the DICE framework's unbiased off-policy learning could inform future work on offline or replay-buffer-based agent training. For researchers building agents that must explore open-ended environments before receiving task specification, SEMDICE provides a theoretically grounded initialization strategy that maximizes state coverage under the worst-case downstream reward.