SEMDICE: Off-policy State Entropy Maximization via Stationary Distribution Correction Estimation
Authors: Jongmin Lee
Meiqi Sun
Pieter Abbeel
Venue: ICLR 2025 Poster (2025)
Keywords: state entropy maximization
unsupervised reinforcement learning
Primary area: reinforcement learning
TL;DR: This paper introduces state-entropy maximization method for RL pre-training based on stationary distribution optimization.

Abstract: In the unsupervised pre-training for reinforcement learning, the agent aims to learn a prior policy for downstream tasks without relying on task-specific reward functions. We focus on state entropy maximization (SEM), where the goal is to learn a policy that maximizes the entropy of the state's stationary distribution. In this paper, we introduce SEMDICE, a principled off-policy algorithm that computes an SEM policy from an arbitrary off-policy dataset, which optimizes the policy directly within the space of stationary distributions. SEMDICE computes a single, stationary Markov state-entropy-maximizing policy from an arbitrary off-policy dataset. Experimental results demonstrate that SEMDICE outperforms baseline algorithms in maximizing state entropy while achieving the best adaptation efficiency for downstream tasks among SEM-based unsupervised RL pre-training methods.

Published as a conference paper at ICLR 2025
SEMDICE: OFF-POLICY STATE ENTROPY MAXI-
MIZATION VIA STATIONARY DISTRIBUTION CORREC- TION ESTIMATION
Jongmin L.e Meiqi Sun Pieter Abbeel-
UC Berkeley YYnssei University
ABSTRACT
In the unsupervised pre-training for reinforcement learning, the agent aims to learn a prior policy for downstream tasks without relying on task-specific reward func- tions. We focus on state entropy maximization (SEM), where the goal is to learn policy that maximizes the entropy of the state's 's stationary distribution. In this paper, we introduce SEMDICE, principled off-policy algorithm that computes an SEM policy from an arbitrary aff-policy dataset, which optimizes the policy di- rectly within the space of stationary distributions. SEMDICE computes single, stationary Markov vate--ntrroyy-maxiiiiii
policy from an arbitrary off-policy dataset. Experimental results demonstrate that SEMDICE outperforms baseline algorithms in maximizing state entropy while achieving the best adaptation ef- ficiency for downstream tasks among SEM-based unsupervised RL pre-training methods.
INTRODUCTION
The essence of intelligent agents lies in their ability to learn from experiences. Reinforcement learn- ing (RL) (Sutton & Barto 1998) provides framework for autonomously acquiring an intelligent behavior by interacting with the environment while receiving reward signals, and it has shown great promise in various domains such as complex games (Mnih et al. [015] Silver et al. (017) and robotic control (Lillicrap et al. 2016)) Haarnoja et al.) 2018). Still, standard RL algorithms learn tabula rasa, i.e. learning from scratch for every task maximize extrinsic rewards without using previously learned knowledge. Consequently, the learned RL policy tends to be brittle and lacks generalization capabilities, limiting its widespread adoption to many real-world sequential decision- making problems (Cobbe et al. 2020 Gleave et al. 2020).
In contrast, in the fields of computer vision (He et al. 2020 Chen et al. 2020) and natural language
processing (Devlin et al. 2019 Brown et al. 2020), large-scale unsupervised pre-training has paved the way for sample-efficient few-shot adaptation. During unsupervised pre-training, the models trained using a large unlabelled dataset, and then the pre-trained models are fine-tuned later for each specific downstream task with handful of task-specific labeled datasets. We consider the analogy setting of the unsupervised pre-training for RL (Laskin et al. (02T). Specifically, during unsupervised RL pre-training, the agent is allowed train over long periods without access to the extrinsic rewards of the environment. This procedure yields pre-trained policy snapshot, aiming for data-efficient adaptation subsequent downstream tasks defined by reward functions that were inaccessible priori.
We consider the State-Entropy Maximization (SEM) approach for RL policy pre-training (Hazan let al.) 2019 Liu & Abbeel 2021b Yarats et al. (021). as it is simple yet provides robust pol- icy initialization for efficient adaptation against the worst-case reward function (Eysenbach et al.) 2022). However, despite its conceptual simplicity and widespread popularity, principled approach to sample-efficient off-policy SEM methods remains unexplored. Existing methods are either on- policy (thus sample-ineffeiiet)) (Hazan al. 2019). off- off-policy but biased (Liu & Abbeel 2021b They construct the intrinsic reward function based on the particle-based entropy
Yarats et al. 2021
"Equal contribution

---

Published as a conference paper at ICLR 2025
estimator (Singh et al., 2003), and then optimize the policy in the direction of maximizing the con- structed intrinsic rewards. For off policy learning algorithms, state particles from the replay buffer naively used to estimate state entropy (Liu & Abbeel 2021b: Yarats et al. (02I), resulting in estimates of the entropy of the state data stored in the replay buffer, rather than the state entropy of the target policy's s state stationary distribution. Consequently, it remains unclear whether these off- policy methods can converge to an optimal SEM policy. One can consider importance sampling to correct the all-pooicyynes (Mutti et al. 202I), but it would suffer from high variance issue due the curse of horizon (Liu et al. (018)) 8), significantly limiting the degree of sample reuse. Consequently, none of the existing methods serves as an unbiased, sample-efficiee SEM method.
In this paper, we present a state-entropy maximization method that precisely addresses the afore- mentioned issues of the existing methods. Our method, Shute-Entrpy-Maxxiimaa DIstribution Correction Extimation (SEMDICE), essentially optimizes in the space of stationary distributions, rather than in the space of policies or Q-functions, and leverages arbitrary off-po samples. We show that SEMDICE only requires solving a single convex minimization problem, and thus it can be optimized stably. To the best of our knowledge, SEMDICE is the first principled and practical algorithm that computes a SEM policy from arbitrary off-policy dataset. In the exper- iments, we demonstrate that SEMDICE converges to an optimal SEM policy in the tabular MDP experiments. Also, regarding RL policy pre-training, SEMDICE adapts downstream tasks more efficiently than eixsting data-based (i.e., SEM-based) unsupervised RL methods (Liu & Abbeel] 2021b Yarats et al. )021).
via stationary
PRELIMINARIES
2.1 MARKOV DECISION PROCESS (MDP)
We assume the environment modeled as an infinite-horizor !1 Markov Decision Process (MDP) M = (S, A, T, r. 7: Po). where S is the set of states 8, A is the set of actions a, : x A - (S)) is the transition probability, r S - IR is the reward function, E [0, 1] is the discount factor, po E A(S) the initial state distribution. The policy over actions" 2 For a given policy, its state stationary distribution d' (s) and state-action stationary distribution d' (s, are defined as follows:
S + A(A) is a mapping from state distribution
(1 7) Σ
7d
Pr(st = s) if? 1
d= ds((s=
(=0
(1)
lim
Σ Pr(st =) if? =
T-200
T+I
1=0
(1 (1-7) 7)
Σ
7² Pr(st = 4 = a) i < 1,
d' as ==
(2)
lim
Σ Pr(st = 8,0 = a) if? =
T-200
T
where so Po. at (()), and 8441 T(s;, az) for each timestep t 2 d= and d= can be un- derstood as normalized (discounted) occupancy measures of (s, a) and respectively. For brevity, we focus on discounted MDPs (7 < 1) in the main text, but our formulation can be easily general- ized to undiscounted (7 1) settings (see Appendix L). We will use the bar notation (7) to denote the distributions for state s, c.g., (s). The goal of standard RL is to learn a policy that maxi- mizes the expected rewards by interacting with the environment: max, (1 E) D2 7rr(s:) B [r(s)] (d .r).
2.2 UNSUPERVISED RL VIA STATE ENTROPY MAXIMIZATION (SEM)
In the unsupervised pre- mre-training of RL Laskin et al. 2021), the agent is trained by interacting with reward-free MDP (S, A, T, 7, Po) over long periods, where the goal is to learn policy that can
We consider infinite-horizon MDPs in the paper for simplicity, but our formulation can be extended finite-horizon MDPs (Appendix M). We call Makovian policy if depends only on the last state and call it stationary policy if does not depend on timestep.
2

---

Published as a conference paper at ICLR 2025
quickly adapt downstream tasks defined by reward functions that are unknown priori. In this work, we are particularly interested in the approach of maximizing the entropy of state stationary
distribution (Jain et al. 2023 Mutti et al., 2021 2022 Liu & Abbeel 2021b: Yarats et al. 2021 as
an objective for unsupervised pre-training RL:
π == arg max H (s)] Σ de (s)s log d' (s)
(3)
In other words, we aim pre-train a policy whose state visitations cover the entire state space equally well. Intuitively, having such an SEM policy would allow agents to efficiently receive reward signals even for arbitrarily sparse reward functions, facilitating fast adaptation for downstream tasks. More discussions on why learning an SEM policy can be useful for RL pre- pre-training can be found in Appendix A
2.3 EXISTING SEM METHOOSS FOR RL PRE-TRAINING
Existing methods for SEM (Lee et al. 2020 Mutti et al. 2021: Liu & Abbeel 2021b; Yarats
al. 2021) commonly follow the following procedures: (1) construct intrinsic reward functions
r(s) R log d" (s), e.g., by particle-based state-entropy estimation: f(*;) ~ - llg d' (s:) e
log (I|s:- -f 1l2) where si-NN denotes k-nearest neighbor of 8;, and (2) update the policy param- eters via policy gradients in the direction of maximizing f. However, this procedure, in principle, requires on-policy state particles { when estimating the state entropy of the current policy (s), implying that the past experiences cannot be reused naively. Still, for off- policy learning to improve sanple-efficiency, existing methods for SEM pre-training (Liu & Abbeell 202 Yarats al. 202T) naively use (off-policy) state particles in the replay buffer for entropy estimation to construct reward functions, combined with an off-policy RI algorithm. This may not be guaran- teed to converge to an SEM policy as they estimate the replay buffer's 's state entropy, rather than the entropy of the target policy's ' state stationary distribution. Importance sampling off-policy corrections (Mutti et al., 2021 of policy gradients, but it suffers from high variance issue with the curse of horizon (Liu et al., 2018). Consequently, existing methods for SEM pre-training are either biased (off-policy) or sample-inefficient (on-policy). The key challenge in devising a sample-efficient SEM algorithm lies in estimating the entropy of the target policy's stationary state distribution using an arbitrary off-policy dataset.
be adopted for
3 SEMDICE
In this section, we derive our SEM method, satte-Etrooy-Mcziizattiii via stationary DIstribution Correction Estimation (SEMDICE), which computes SEM policy from arbitrary off-policy dataset. Our derivation starts by formulating the regularized SEM problem via concave program- ming that directly optimizes stationary distributions, rather than optimizing policy. All the proofs can be found in Appendix
3.1 CONCAVE PROGRAMMING FOR SEM
Some careful readers may wonder which policy should be considered obtain an optimal SEM policy, as it may not be immediately evident that searching for an optimal SEM policy within sta- tionary Markov policies (instead of richer non-Markovian policies) is sufficient. Fortunately, the following proposition addresses this concern.
Proposition 3.1. (Hazan et al. 2019) There always exists stationary Markovian policy that max- imizes the entropy of state stationary distribution (iee. solution of (3)). Such an optimal stationary Markovian policy is generally stochastic.
By Proposition 3.1 it is sufficient to consider the stationary distribution induced by stationary Markovian policies to obtain an optimal SEM policy. We then begin our derivation with the fol- lowing concave programming problem that optimizes the stationary distributions to solve a (regu- larized) SEM problem, where the primary objective is maximize the entropy of state stationary
3

---

Published as a conference paper at ICLR 2025
distribution of some target policy while adopting f-divergence regularization.
Σ d(s) log d(s) -D:, (d(s, l) (s, a)) max d,d20
(4)
regularizer
state entmpy Fl(d(())
s.t. Σ d(s's a²) = (l)os)p + dss, E ())((ss)) 4s
(5)
Ed(,,o)= d(s) As
(6)
The Bellman flow constraint (5) guarantees that d(s, a) is a valid state-action stationary distribution normalized occupancy measure (s.a). The
for some policy, where d(s, a can be understood marginalization constraint (6) ensures that d(s) is the state stationary distribution directly induced by d(s,a). In 44), Dy(((s, a)l d (s, a)) == (s.a)>a'' [S (xa))))) )] denotes the f-divergence between d and d8, h= = {(x, s');));))))))) d'); denotes any off-policy dataset of interest (e.g., replay buffer of the agent), and dD is its corresponding distributions. We assume f is a strictly convex function and continuously differentiable. For brevity, we will abuse the notation dD to represent (s, ~ dD, (s, a,s) ~ dP. The regularization hyperparameter a > O balances between maximizing state entropy and preventing distribution shift from past experiences.
The regularization term (y(dldd)) (d)) in (4) can be understood as imposing trust-region updates Schul- man et al. (2015) by constraining the solution to the vicinity of previous state-action visitations. Technically, the regularization term ensures strict concavity of the objective function with respect to d. thereby guaranteeing the uniqueness the optimal solution for the concave programming (4. 6). In contrast to existing DICE-based RL methods (Lee et al. 2021 2022, Nachum & Dai 2020) that formulate optimization problems only for state-action stationary distribution, we define the optimization problem that jointly optimizes state stationary distribution with the marginalization constraint (6) to deal with state entropy.
To sum up, seek state-action stationary distribution d (and its corresponding state stationary dis- tribution d) whose state entropy being maximized. Once we have computed the optimal solution (d", dd) of the concave programming 64)) its corresponding optimal policy can be obtained by (s.a)
normalized d' for each state (Puterman 1994): x (als) d=(s)
3.2 LAGRANGE DUAL FORMULATION
Still, solving 4 6 directly requires white-box model of the environment, which is inaccessible in many practical applications of RL. To derive an algorithm that can be fully optimized a model-free manner, we consider the Lagrangian for the constrained optimization problem (4
max min d,d20 K.
) Σd ((s)) log d(s) - aD, + Es() -) d(ss)
Ev((s) (1- + d((, (((ss - E. d(s',a)) a'))
8.0
where v(s) IR are the Lagrange multipliers for the Bellman flow constraints 55. and u(s) E R are the Lagrange multipliers for the marginalization constraints 6). Then, we rearrange the terms in
max min(1 Beo [v(5a)] aLE (s.))oD) [5( ([d(tee) d,020
(8)
d(1 a) ((s) + 7B, (('') u(s)
Σ d(ss) ((((()(s)))))))))))))))()))))( log d(s)
(s.o)((..o)
min max (1 7) Eno
[v(s))] CEE
(( (
+ a(s,,)) d(s)) (n(s) + log d(ss)
b, 4,429
In we could reorder maximin to minimax thanks to strong duality (Boyd et al. 2004). Finally, using the Fenchel conjugate, we can eliminate the inner maximization problem and end up with single convex minimization problemf]
This simplification from min-max to min made possible by introducing the optimization variable d along with the marginalization constraint @ (see Appendix
4

---

Published as a conference paper at ICLR 2025
Theorem 3.2. The minimax optimization problem 9 is equivalent solving the following uncon- strained mininization problem:
min - Ep p + E o:f (d + exp(((-((s) - =: L(v, p))
(10)
where (z) maXx20 zy (z)) The objective function L(v, (4) is convex for v and Fur- thermore, given the optimal solution (V","') = arg min L(v, pi), the stationary distribution corrections of the optimal SEM policy are given by: (s)).))
d (1 (s(())
7B (V (8') -1))) =: =: =. (s,
(5))(s
(11)
EE (8,0)
where x4 20)
In short, by operating in the space of stationary distributions, computing a stationary Markovian SEM policy can, in principle, be addressed by solving convex minimization problem. The resulting "Cn,," (k,a)d (e.0) in the case of finite policy can be obtained by N² (a Ed d' (8,0²) En wit "er wu- (8,0) ( (s,ar) MDPs.
(8.0))
3.3 PRACTICAL ALGORITHM
Still, one practical issue in 100) is that optimizing it is unstable due to its inclusion of exp(-) term, often causing exploding gradient problems. To remedy this issue, consider the follow- ing numerically stable alternative objective function L that replaces E, expp(=p(e - 1) with log E, exp(-M(s)).
Theorem 3.3. Define the objective functions L(v, p)) as
L) E) [M(s))] (s.aaaa [ast (tena(s, + log log exp(p11p m(s)) (12)
Then, for any optimal solutions (v" ,p²) = arg min,, L(v, H) and (V", ,T*) = arg min [(v p), the following holds:
I(v",p") L(V",''') ,T")
(13)
Also, there exists constant C such that the following holds:
μ¹ = P + C and v" "0' +
(14)
exp(=u(s)) Note that the gradient V, E, log exp((x(((s) - E, exp(-p(>)) Vsp(s) normalizes exp(-)
by softmax, thus L no longer suffer from numerical instability by large gradients. At first glance, it seems that optimizing 12) could yield a solution that is completely different from that of (10). However, Theorem 3.3 shows that their optimal objective function values are the same and their optimal solutions only differ in constant shift. Furthermore, despite its constant shift, it does not change anything for w" computation in (II). as can be seen as follows:
Ei f (s, a) = (1²(s) C) + (E(P"(s') ( (v'(s) C
= 1r" (x) + 7B - - (s) (8)= ev- u (8,a)
w (8, a) ww. p ())a)
Still, (12) requires summing up (or integrating) every possible state. which is intractable for large (or continuous) state space. Therefore, we approximate the log Σ exp( (((()))()) term via Monte-Can integration with arbitrary distribution q: log Σ exp(-nu(s) Finally, optimize the following objective function with g(s) (s):
log (x((e(((=u(s) log (s)))
min(1 ?) Eso [s(s)]] + ( [os (eeu (8:
(15)
log B.ndd(s) [exp(=p(s) - log d (s))
5

---

Published as a conference paper at ICLR 2025
The last remaining challenge is that (5) still cannot be naively optimized in fully model-free manner in continuous domains, as it contains computing expectations w.r.t the transition model inside the non-linear function f+ (.). For a practical implementation, we use the following objective that can be easily optimized via sampling only from dD:
min(1 7) () + (' a14 (de Kuje
(16)
+ log E (o)) exp( (xp(-m(s) log d0 (s)) (.(,
where a') := =us)) + 7p((s) v(s) is single-sample estimate of e(s, a). Note that every term in 166 can now be evaluated only using the (s, a. s') samples from D, thus we can optimize y and μ easily. Although (16) is a biased estimate of (15). one can easily show that [(v, μ) is an upper bound of L(v, ji), i.e. ((,, I L(v, (4) always holds, where the equality holds when the MDP is deterministie, by Jensen's 's inequality and the convexity of / (-). Once we get the optimal solution (E', ""') = arg min [(,, pi), we end up with the stationary distribution corrections of the optimal SEM policy: (s,a)
(P.0)) wa- f (s, a) by (11 In practice, to deal with continuous state space, we parameterize vo : S - R and nr : S - R using simple MLPs, which take the state as an input and output scalar value, and optimize the network parameters.
Policy Extraction The final remaining question is how to obtain an explicit policy, our method computes the distribution correction ratios w" instead of the policy itself directly. We follow the i-projection used in (Lee et al. 2021).
min KL (s)(ss))
(17)
F Ba
[log w" (s, (KL(((()s)] +
(18)
Essentially, minimizing 188) corresponds to computing policy (-Is)) that mimics the optimal SEM policy (-(ss) for each state s E D. In summary, SEMDICE computes the SEM policy as follows: First, minimize (I6), and Second, extract the policy using the obtained Wy-. by i-projection (88)
In practice, rather than training until convergence at each iteration, we perform a single gradient update for v, μ. and æ. We outline the details of policy extraction (Appendix) FF and the full learning procedure with pseudo-code in Algorithm
( (Appendix (
Remark 3.4. Minimizing L(e, (i) (5) equivalent to solving the original (regularized) state- catroyy-maxmizatton problem (4 demonstrating that computing (regularized) SEM policy can, in principle. optimizes the state entropy of target policy, rather than optimizing the state entropy of the replay buffer as in (Yarats et al. [2021] Liu & Abbeell 2021b). To the best of our knowledge, SEMDICE is the first principled and practical off-policy SEM algorithm that can learn from arbitrary off off-policy experiences. Also, (Hazan et al.] 2019) showed that state entropy is not concave for policy param- eters, whereas state entropy concave for the stationary distribution space d(s). That being said, our method, which directly optimizes stationary distributions, may offer better convergence prop- erties compared to existing policy- based algorithms. However, providing formal analysis for the convergence guarantee remains as future work.
achieved by arbitrary off-policy dataset D. Note that this approach directly
4 RELATED WORK
Unsupervised RL and State Entropy Maximization Our work falls within the realm of unsu- pervised reinforcement learning, specifically addressing the challenge of task-agnostic exploration. The reward-fiee freard-free exploration framework has gained significant attention in recent years (Jin et al.) 2020 Tarbouriech et al. [2020] Kaufmann et al. 2020). While these methods share a similar context to our work, they pursue largely orthogonal objectives.
State entropy maximization is a particular instance of reward- exploration objective, where the agent aims to estimate the density of states and maximize entropy (Hazan et al. 2019 Lee et al.
2020
Liu & Abbeel 20216; Yarats et al.] 202T; Mutti et al.) 2022] Tiapkin et al. 2023 Kim et al.
2023
Yang & Spaan 2023 However, these methods mostly rely policy gradients with the

---

Published as a conference paper at ICLR 2025
Uniform
Optimal
CB-SA
CB-S
PB-S
SEMDICE (ours)
Entropy of (s)))
Entropy of d21s)
1.0
normalized entropy
normalized entropy
0.5
0.5
0.0
10p the number episodes
the number episodes
Figure I: Performance of SEMDICE and baselines in randomly generated tabular MDPs. The first column indicates normalized policy entropy (0: state entropy of uniform policy, 1: state entropy of optimal SEM policy), the second column indicates normalized entropy of the cumulative experi- ences. We observe that only SEMDICE converges an optimal SEM policy, whereas baselines do not directly maximize the state entropy.
reward function for the estimated state entropy, and they are either sample inefficient due to their requirements of on-policy samples (Hazan et al.) 2019 Mutti et al., 2022), optimize the biased objective due to estimating the state entropy of the replay buffer (Liu & Abbeel 202 lb Yarats al. 2021). In contrast, our method operates directly in the space of stationary distributions and serves as the first principled off-policy SEM method.
Besides state entropy maximization. various self-supervised objectives have been explored for unsu- pervised RL pre-training (Laskin et al. 202T). The goal of pre-training in this context is to compute policy without relying task-specific reward functions, enabling rapid adaptation future, un- known downstream tasks defined by reward functions. Unsupervised RL pre-training algorithms are categorized three main ways (Laskin al.) 2021). Knowledge-basee methods aim increase knowledge about the environment by maximizing prediction error (Pathak et al. 2017, [Burda et al, 2019 Pathak al. 2019). Compefence- based methods (Lee et al., 2020] Eysenbach et al. 2019 Liu & Abbeel [02 |bbeel][02laa Laskin et al.) (022) aim to learn explicit skill representations by maximizing mutual information between encoded observation and skill. Lastly, data-based methods (Liu & Abbeel [021b] Yarats et al., (02T)) aim to achieve data diversity via particle-based entropy maxi- mization. Our SEMDICE is categorized into the data-based method when considering unsupervised
RL pre-training.
Stationary DIstribution Correction Estimation (DICE) DICE-family algorithms perform sta- tionary distribution estimation and have shown significant promise various off-policy learning sce- narios in reinforcement learning (RL), such off-policy evaluation (Nachum et al., [019a] Zhang Zhang et al. 2200 Yang et al. 2020), reinforcement learning (Lee et al. 202 1| Kim
et al. 2020
et al. 2024
Mao et al. 2024), constrained RL (Lee et al. 2022), imitation learning (Kim et al.)
2022b Ma et al. |l.]| [022] Kim et al. 2022a Sikchi et al. 2024), and more. However, to the best of our
knowledge, DICE-based algorithms have been proposed to address state entropy maximization.
EXPERIMENTS
In this section, we empirically evaluate our SEMDICE and baselines: to demonstrate (1) SEMDICE converge to an optimal SEM policy, (2) visualization of SEMDICE's state visitation, (3) how effi ciently SEMDICE pre-trained policy can adapt to downstream tasks on URL benchmarks (Laskin) al. 2) Ablation experiments on the choice of f and for SEMDICE can be found in the
Appendix K
5.1 STATE ENTROPY MAXIMIZATION IN FINITE MDPs
We first evaluate SEMDICE and baseline algorithms on randomly generated finite MDPs with 20 states and actions to show how effectively SEMDICE maximizes the state entropy. compared to baselines, when optimized using the off-policy dataset (the entire replay buffer). We conduct repeated experiments with different seeds for 100 runs.

---

Published as a conference paper at ICLR 2025
Baselines Baselines are rolicy-gradient-bee based methods, with different (non-stationary) intrinsic reward functions A. We consider the following baselines, and they can be thought of as analogies to existing unsupervised RL approaches. (1) CB-SA: it constructs the state-action-ccuutt-asss ex- VN(s.aj" where N(s, a) is the cumulative (s, a) ()--isitatiin counts
ploration rewards by F(s, a)
by the agent. This baseline can be understood an analogy to the existing count-based methods or RND (Burda et al. ) (2) CB-S: it constructs sate-counttbaas baaunt-based rewards f(s)
VN(a)* where
(s) is the s-visitation counts. It is similar to CB-SA but only considers state visitation frequencies. thus is expected to be closer to state-entropy maximization. (3) PB-S: it constructs entropy-bant log (s), where dD is the empirical state distribution of the cumulative experi-
rewards f(s)
ences (which is not direct estimate of log d (s)). This baseline can be understood as an analogy
to existing SEM methods performing off off-policy updates (Liu & Abbeel 2021b Yarats et al. 202
(4) Uniform: sample actions uniformly at random.
For all baselines, we first collect trajectories using the current policies and then construct the non-stationary rewards function based on visitation Lastly, we perform policy gradient ascents along the direction of maximizing the constructed reward function.
measures accordingly.
Evalution For each run, we (1) generate random MDP and initialize the policy uniform policy; (2) For each method, use the current policy π to collect 10 episodes and add them to the replay buffer D; (3) Using D, compute the MLE transition matrix Т, the intrinsic reward function A, and the value function O with respect the intrinsic reward; (4) Perform aff-policy policy gradient ascent based on the estimated O (5) Repeat steps (2)-(4) until D contains 1000 episodes. Throughout the process, we compute the state entropy of the current policy d= (s) as well as the entropy of the replay buffer dd (s).
Results Figure I presents the results. We observe that CB-S baseline generates slightly better SEM policy than CB-SA, which is natural as it focuses on the exploration of the unvisited states while ignoring action uncertainty. PB-$ baseline could obtain nearly maximal state-entropy for the dataset in the dataset ds).). but it is not directly optimizing the target policy's state entropy d (s). SEMDICE is the only algorithm that could converge to an optimal SEM policy efficiently. achieves much better sample efficiency in terms of the state entropy of dataset that PB-S. Although we didn't 't include the result, we also tested on-policy policy- gradient approach for SEM. It could converge to optimal SEM policy with carefully chosen learning rate, but it required x 100 more samples to reach near-optimal entropy maximization. This results confirm that SEMDICE is capable of computing an optimal SEM policy even from the arbitrary off-policy dataset. In Appendix [ further show that SEMDICE still converges to the optimal SEM policy even when the dataset is collected by purely off- off-policy agent (i.e. uniform random policy).
also
5.2 STATE ENTROPY MAXIMIZATION VISUALIZATION
This section aims to visualize the behavior of SEMDICE and baselines for their state visitations during policy learning. To this end, we use MountainCar (Continuous) from Gymnasium (Towers) et al. 2023), classic control task with low-dimensional continuous state and action space. The task involves a deterministic MDP where a car must escape from sinusoidal valley by building momentum. The observation space is two-dimensional, consisting of position E [-1.2, 0.6 and velocity E [-0.07, 0.07], while actions correspond to applying directional forces in the range [-1, 1]. The low dimensionality allows for clear visualization of state visitation distributions in the 2D plane, facilitating straightforward comparisons between different agents.
Baselines We compare SEMDICE to baselines from three categories of unsupervised RL. Knowledge-based methods include RND (Burda et al. 2019), which relies on prediction error as an exploration signal. Data-based methods aim to directly maximize state entropy, which can be our direct interest, and include ProtoRL (Yarats et al.) (021). APT (Liu & Abbeel| 2021b). and MEPOL
Mutti et al. 202T). Competence-based methods such as DIAYN (Eysenbach et al. (019)) and CIC
Laskin et al. 2022) focus on learning diverse skills by maximizing the mutual information between skills and observations.
8

---

Published as a conference paper at ICLR 2025
SEMIDICE (oun)
PROTO
APT
MEPOL
CIC
DIAYN
RND
Snapchot
velocit
Entropy 1:24
Esnaay 5.27
Entree 1
Entrepy 4-48
Estraay:444
Entrops 1:34
Ean
Snapshot 100K
veloctty
Emtray
Entrayy 1.41
Estreay A
5:00
200K
-200
.55 posibion
possion
posibion
position
Figure 2: Mountaincar state coverage: Visualization of agent pretraining results across SEMDICE and baselines. During the reward-free pretraining stage, we save model snapshots at 10K, 100K, and 200K steps. From each snapshot, draw 30k samples via environment interactions for visu- alizations. The state space is discretized into uniform 51 x 51 bins, and the empirical distribution normalized to compute the entropy shown the bottom of each subplot. While existing methods failed to converge to a cat--entrroy-mximmiiii policy, SEMDICE successfully did so.
SEMDICE (ours)
PROTO
APT
CIC
Disagreement
RND
MEPOL
Walker
Quadruped
Jaco
entropy index
60
entropy Index
entropy Index
-40
40
100
20
-60
0.1
0.3 0.5 Iterations
1.0
2.0
0.1
0.3 0.5 Iterations
1.0
2,0
0.1
0.3 0.5 Iterations
1.0
2.0 1e6
le6
1e6
Figure 3: Particle-based entropy estimations of pretrained SEMDICE and baselines in URLB (Laskin et al. 2021). We follow the setup from URLB and pretrain each agent 2M environment interactions. We save policy snapshots at [0.1M, 0.3M, 0.5M, 1M, 2M] steps, and report the change of entropy index along the train- ing process. We observe that SEMDICE achieves maximum state entropy in both Walker and Jaco Arm do- main, and close maximum state entropy Quadruped domain. It is also the most sample efficient approach amongst all agents: with few 100K environment interactions, SEMDICE learns proper state entropy maximizing policy. For example, in Walker domain, multiple agents like RND, Disagreement, and PROTO managed to achieve similar state entropy as SEMDICE. However, they are much less sample efficient, and only converged to maximum state entropy policy after close to 2.M environment interactions.
Evaluation and Results We evaluate all methods based on their overall state coverage and entropy during the reward-free pretraining phase. Specifically, take snapshots of each policy at 10K, 100K, and 200K environment steps under purely intrinsic rewards. For each snapshot, we collect 30,000 state samples by interacting with the environment using the fixed policy, and visualize the corresponding state visitation distribution. To compute the entropy of the stationary state distribution induced by each policy, we discretize the 2D continuous observation space into 51 51 grid and count the number of visits to each grid cell. Note that all methods operate directly in continuous state and action spaces; discretization used solely for entropy computation.
Figure 2 shows that existing state entropy maximization methods (PROTO, APT, MEPOL) fail to converge single SEM policy. Instead, their learned policies exhibit non-stationary behavior throughout training, likely due to their dependence on non-stationary, learned reward function. In contrast, SEMDICE progressively increases state coverage over time and eventually converges to single, stable SEM policy, which is in contrast to the oscillatory behaviors observed in the baselines.
5.3 URL BENCHMARK
Finally, evaluate the state-entropy maximization and fine-tuning RL performance using tasks from URLB (Laskin et al. 202T). URLB consists of twelve tasks across three different domains: Walker, Quadruped, and Jaco Arm. During the pre-training phase, SEMDICE and baselines com-

---

Published as a conference paper at ICLR 2025
Modian
Mean
Optimality Gap
SEMDICE
Data
PROTO
APT
Competence
DIAYN
Knowledge
DISAGREEMENT
RND
0.15 0.30
0.45
0.15
0.30
0.45
0.15
0.30
0.45
0.60
0.75
Expert Normalized Score
Figure 4: We evaluate all methods on the 12 tasks provided by the URLB benchmark (Laskin et al. 2021). covering three domains: Walker (stand, walk, flip, run) and Quadruped (stand, walk, run, jump) are locomotion tasks that require maintaining balance and strategically controlling the body; Jaco Arm (reach bottom left, bottom right, top left, top right) consists of manipulation tasks involving precise control of a 6-DOF robotic arm with a three-finger gripper to reach specified target positions. To ensure statistically reliable comparisons, we rliable (Agarwal al. 202I) to report aggregate statistics across multiple seeds and tasks, along with stratified bootstrap confidence intervals. Each plot is generated using total of 120 runs (3 domains x 4 tasks/domain 10 seeds).
putes the pre-trained policy without task-specific reward, solely relying on their intrinsic objective functions. Once the policy pre-training is done after 2M steps, a small amount of environment in- teractions are allowed with the task-specific rewards (100K steps). We then evaluate how efficiently each method adapts to the given downstream tasks, specifically examining how well the pre-tra pre-trained policies serve as a good initialization for rapid adaptation. We compared SEMDICE with the base- lines provided by URLB. We follow the experimental protocoal of URLB: ran 10 seeds per task, and used the same hyperparameters for baselines, except for using the smaller MLP hidden size 1024 256 for fast training/evaluation.
Results In Figure 3 we report the state entropy of the policy checkpoints during pre-training per method, and SEMDICE shows the highest state entropy estimates across all domains. Fig- ure 4. we present the fine-tuning performance. Overall, SEMDICE significantly outperforms all data-based (i.e. SEM-based) and knowledge-basee baselines in terms of rapid adaptation perfor- mance during fine-tuning phase. This highlights that state entropy maximization is a desirable ob- jective for RL pre-training, as illustrated in Appendix(A] Although SEMDICE underperforms CIC, a competence-ber based URL method, we believe its performance can be further enhanced by integrating representation learning techniques (e.g.. those used in CIC). is important to emphasize that our results were achieved solely through the SEM objective without additional mechanisms. Moreover, SEMDICE's potentially more unbiased and efficient SEM optimization could have resulted in good policy initialization for fine-tuning. Qualitatively, we observed that even the data-based baselines sometimes exhibit static behavior, while SEMDICE's learned policy consistently showed actively moving behavior across different domains.
CONCLUSION
We presented SEMDICE, new stat-entroyy-mxxiiiaatoo (SEMDICE) algorithm for RL pre- training. Existing SEM methods rely on policy gradients, and they either require on-policy samples or perform (off-policy) biased optimization due to its estimation of state entropy of replay buffer. In contrast, SEMDICE directly optimizes in the space of stationary distributions, and our derivation shows that computing an optimal SEM policy can be achieved by solving single convex minimiza- tion problem with arbitrary off off-policy dataset. To the bast of our knowledge, SEMDICE is the first principled off-policy SEM algorithm. Through various tabular and continuous domain experiments, demonstrated that SEMDICE converges to single stationary Markov SEM policy, and outper- forms baselines in terms of maximizing state entropy. As for future work, we plan to incorporate representation learning components into SEMDICE so that it enables faster meaningful feature de- tection and better navigation in high-dimensional domains such as pixel-based domains. Another promising future direction is to devise DiCE-based method for competence-be unsupervised RL algorithms (e.g. maximizing the mutual information), which can serve as principled off off-policy algorithm for
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported by NSF AI4OPT AI Centre. Pieter Abbeel holds concurrent appointments as a Professor UC Berkeley and as an Amazon Scholar. This paper describes work performed at UC Berkeley and is not associated with Amazon. This work was partly supported by Institute of In- formation & communications Technology Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT) (No. RS-2024-00457882, Al Research Hub Project) and the Institute of In- formation & Communications Technology Planning & Evaluation (ITPP) grant (RS-2020- 122001120131,, Artificial Intelligence Graduate School Program (Yonsci University).
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. In Advances in Neural Infor- mation Processing Systems, volume 34, pp. 29304-29320.. Curran Associates, Inc., 2021.
Stephen Boyd, Stephen P Boyd, and Lieven Vandenberghe. Convex optimization. Cambridge uni- versity press, 2004.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901, 2020.
Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. In International Conference on Learning Representations, 2019.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1597- 1607. PMLR, 13-18 Jul 2020.
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 2048-2056. PMLR, 13-18 Jul 2020.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Pupers), pp. 4171-41866 Association for Computational Linguistics, June 2019.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity all you need: Learning skills without reward function. In International Conference on Learning Representa- tions, 2019.
Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. The information geometry of unsupervised reinforeement learning. In International Conference on Learning Representations, 2022.
Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized markov decision processes, 2019. URL httpps:s:///axiiie org/ abs/ /99//1911.01 11275
Adam Gleave, Michael Dennis, Cody Wild, Neel Kant, Sergey Levine, and Stuart Russell. Adver- sarial policies: Attacking deep reinforcement learning. In International Conference on Learning Representations, 2020.
11

---

Published as a conference paper at ICLR 2025
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference Machine Learning (ICML). 2018.
Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum en- tropy exploration. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2681-2691. PMLR, 09-15 Jun 2019.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xic, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9726-9735 2020.
Arnay Kumar Jain, Lucas Lehnert, Irina Rish, and Glen Berseth. Maximum state entropy explo- ration using predecessor and successor representations. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. Reward-free exploration for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 4870-4879. PMLR, 13-18 Jul 2020.
Emilie Kaufmann, Pierre Ménard, Omar Darwiche Domingues, Anders Jonsson, Edouard Leurent, and Michal Valko. Adaptive reward-free exploration, 2020. URL https::/arxiv..rgg abs/2005 06294
Dongyoung Kim, Jinwoo Shin, Pieter Abbeel, and Younggyo Seo. Accelerating reinforcement learning with value-conditional state entropy exploration. In sehirty-seennth Conference on Neural Information Processing Systems, 2023.
Geon-Hyeong Kim, Jongmin Lee, Youngsoo Jang, Hongseok Yang, and Kee-Eung Kim. LobsDICE: Offline learning from observation via stationary distribution correction estimation. In Advances in Neural Information Processing Systems, 2022a.
Hron-Hyeong Kim, Seokin Seo, Jongmin Lee, Wonseok Jeon, HyeongJoo Hwang, Hongseok Yang, and Kee-Eung Kim. DemoDICE: Offline imitation learning with supplementary imperfect demon- strations. In International Conference on Learning Representations, 2022b.
Woosung Kim, Donghyeon Ki, and Byung-Jun Lee. Relaxed stationary distribution correction esti- mation for improved offline policy optimization. Proceedings the AAAI Conference Artificial Intelligence, 38( 12): 1111111111119111 Mar. 2024.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. URLB: Unsupervised reinforcement learning benchmark. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Michael Laskin, Hao Liu, Xue Bin Peng. Denis Yarats, Aravind Rajeswaran, and Pieter Abbeel. Unsupervised reinforcement learning with contrastive intrinsic control. In Advances in Neural Information Processing Systems, 2022.
Jongmin Lee, Wonseok Jeon, Byung-Jun Lee, Joelle Pineau, and Kee-Eung Kim. Optidice: Offline policy optimization via stationary distribution correction estimation. In Proceedings of the 38th International Conference on Machine Learning (ICML). 2021.
Jongmin Lee, Cosmin Paduraru, Daniel J Mankowitz, Nicolas Heess, Doina Precup, Kee-Eung Kim, and Arthur Guez. COptiDICE: Offline constrained reinforcement learning via stationary distribu- tion correction estimation. In International Conference on Learning Representations, 2022.
Lisa Lee, Benjain Eysenbach, Emilio Parisotto, Erix Xing, Sergey Levine, and Ruslan Salakhutdi- nov. Efficient exploration via state marginal matching, 2020.
12

---

Published as a conference paper at ICLR 2025
Timothy P. Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In 4th International Conference Learning Representations, ICLR, 2016.
Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 6736-6747, 18 24 Jul 2021a
Hao Liu and Picter Abbeel. Behavior from the void: Unsupervised active pre-training. In Advances in Neural Information Processing Systems, 2021b.
Qiang Liu, Lihong Li, Ziyang Tang, and Dengyong Zhou. Breaking the curse of horizon: Infinite- horizon off-policy estimation. Advances neural information processing systems, 31, 2018.
Yecheng Jason Ma, Andrew Shen, Dinesh Jayaraman, and Osbert Bastani. Versatile offline imitation from observations and examples via regularized state-occupancy matching, 2022. URL https:
2ttp:/2222.022222 02433
Liyuan Mao, Haoran Xu, Weinan Zhang, and Xianyuan Zhan. ODICE: Revealing the mystery of distribution correction estimation via rrthogonal-gradien update. In The Twelfth International Conference on Learning Representations, 2024.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, (29 533, 2015.
Mirco Mutti, Lorenzo Pratissoli, and Marcello Restelli. Task-agnostic exploration via policy gra- dient of a non-parametric state entropy estimate. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 9028-9036, 2021.
Mirco Mutti, Riccardo De Santi, and Marcello Restelli. The importance of non-markovianity in maximum state entropy exploration. In International Conference on Machine Learning, pp- 16223-1622 PMLR, 2022.
Ofir Nachum and Bo Dai. Reinforcement learning via fenchel-rockafellar duality, 2020.
Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. DualDICE: Efficient estimation of off off-py stationary distribution corrections, 2019a.
Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and Dale Schuurmans. Algaedice: Policy gradient from arbitrary experience, 2019b. URL httlhtts///axiiv.rrtt http://arxi...or//////1122 02074
Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervise prediction. In Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 2778-2781 06-11 Aug 2017.
Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta. Self-supervised exploration via disagreement. In Proceedings of the 36th International Conference Machine Learning, volume 97 of Pro- ceedings of Machine Learning Research, pp. 5062-5071, 09-15 Jun 2019.
Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. John Wiley & Sons, Inc., USA, 1st edition, 1994. ISBN 0471619779.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1889-1897, Lille, France, 07-09 Jul 2015. PMLR.
13

---

Published as a conference paper at ICLR 2025
Harshit Sikchi, Qinqing Zheng, Amy Zhang, and Scott Niekum. Dual rl: Unification and new methods for reinforcement and imitation learning, 2024. URL atttp::/arxiv..om abs/
2332.08566
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap, Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of Go without human knowledge. Nature, 550:354-359, 2017.
Harshinder Singh, Neeraj Misra, Vladimir Hnizdo, Adam Fedorowicz, and Eugene Demchuk. Near- est neighbor estimates of entropy. American Journal of Mathematical and Management Sciences, 23:301 321, 2003.
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: An introduction. MIT Press, 1998.
Jean Tarbouriech, Shubhanshu Shekhar, Mattco Pirotta, Mohammad Ghavamzadeh, and Alessandro Lazaric. Active model estimation in markov decision processes, 2020. UtLptttss:/aaxxv
org/ 20c//bb/220 abs/2003. 03297
Daniil Tiapkin, Denis Belomestny, Daniele Calandriello, Eric Moulines, Remi Munos, Alexey Nau- mov, Pierre Perrault, Yunhao Tang, Michal Valko, and Pierre Menard. Fast rates for maximum entropy exploration. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 34161-34221, PMLR, 23-29 Jul 2023.
Mark Towers, Jordan K. Terry, Ariel Kwiatkowski, John U. Balis, Gianluca de Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Arjun KG, Markus Krimmel, Rodrigo Perez- Vicente, An- drea Pierré, Sander Schulhoff, Jun Jet Tai, Andrew Tan Jin Shen, and Omar G. Younis. Gymma- sium, March 2023.
Mengjiao Yang, Ofir Nachum, Bo Dai, Lihong Li, and Dale Schuurmans. Off-policy evaluation via the regularized lagrangian. In Advances in Neural Information Processing Systems, volume 33, pp. 6551 6551-6561. Curran Associates, Inc., 2020.
Qisong Yang and Matthijs T.J. Spaan. Cem: Constrained entropy maximization for task-agnostic safe exploration. Proceedings of the AAAI Conference on Artificial Intelligence, 37(9):10798- 10806, Jun. 2023.
Denis Yarats, Rob Fergus, Alessandro Lazaric and Lerrel Pinto. Reinforcement learning with proto- typical representations. In Proceedings of the 38th International Conference on Machine Learn- ing, volume 139 Proceedings of Machine Learning Research, pp. 1 11920-11931. PMLR, 18-24 Jul 2021.
Ruiyi Zhang* Bo Dai*, Lihong Li, and Dale Schuurmans. Gendice: Generalized offline estimation of stationary values. In International Conference on Learning Representations, 2020.
Shangtong Zhang, Bo Liu, and Shimon Whiteson. Gradientdice: Rethinking generalized offline estimation of stationary values. In Proceedings of the 37th International Conference on Machine Learning, ICML. 2020, 13-18 July 2020, Virtual Event, volume 1 119 of Proceedings of Machine Learning Research, pp. 11194-1 1203. PMLR, 2020.
14

---

Published as a conference paper at ICLR 2025
A WHY STATE-ENTROPY MAXIMIZATION FOR RL PRE-TRAINING
In this section, we illustrate why state-entropy maximization can be good choice for RI. pre- training, Consider the following information anformation-theoreti theoretic regularized regret objective defined by (E--
senbach et al. 2022).
DOAPAATIONOBBECCIVE(C.....
min max
(1)))))) +Dxi (d' (d)
(19)
In the ADAPTATIONOBJECTIVE,, for given pre-trained policy's stationary distribution dpre and downstream reward function r, it measures the regret of the policy (d²) after adaptation from the pre-trained policy (dpre) with the information cost, where optimal dt (s) denotes the state stationary distribution of the optimal policy for the reward r (introduced define regret). Then, a desirable pre-trained policy (or its corresponding d²)) can be defined in terms of minimizing the (regularized) regret against the worst-case reward function:
min max ADAPPATIONOORICTIVVV (8), r())
(20)
4pre(s) r(s)
Eysenbach al. 2022 has shown that (20 equivalent to:
min max (d (s)[[apre(())
(21)
(pre(s) d=)))
and thus maximum-entrropy (uniform) dpre is the solution of 21 and (0). This means that SEM policy provides robust policy initialization against the worst-case reward function for fine-tuning,
justifying why state- state-entrppy maximization is useful for unsupervised RL pre-training. Figure
visualizes information geometry of RL pre-training (Eysenbach 2022
d (s2))
Reward
d(s,)
(
Figure 5: Visualization reward function and policies as their state stationary distributions 3-state MDP. The shaded orange denotes of achievable state stationary distributions that lie on the probability simplex (SS) (shaded blue). The green arrow represents the reward function vector starting the origin. denotes the stationary distribution of pretrained SEM policy, dpre = 33 3 3/. and denotes the state stationary distribution of an optimal policy for r, dt which is an intersection of and A(S). The distance from
reflects the adaptation cost.
15

---

Published as a conference paper at ICLR 2025
B
WHY BOTH d( AND d(s) ARE USED
Without d(s), the main objective becomes: Ed(s,oo log Edss,a')
(22)
This makes algorithm derivation complicated due to the summation inside logarithm. Specifically, our derivation with d(s) results in the following inner-max problem (See (23) in Appendix C
max Eg (rsss a) (p(s)) (msss(())))))))))) - -ff w=2
(w (s)) a))] -af) (s,
af; (: (p(s) (
where inner-maximization problem maxw20 could have been easily eliminated by applying Fenchel conjugate of f.e. f+ (v) max;20 ry fz)
In contrast, without dsss, we have the following inner-maximization problem:
max ((,a))ad (ss a) ( log Σ w(s, a' d'd (s, (s, + [v('] v(s)) - at ( #20
where Fenchel conjugate directly applicable eliminate this inner-maximization problem.
To sum up, the introduction of the new optimization variable d(s) eliminates the need for solving nested min-max optimization, enabling SEM be formulated solving single convex optimiza- tion problem.
C PROOFS
C.I PROOF OF PROPOSITION 3.1
So
$2
a
Figure 6: An illustrative example of MDP with three states and two actions, where the optimal SEM policy is stochastic. Specifically, in order to maximize the state entropy, the action selection should be randomized s1. If the agent deterministically takes action a 0 at $1, it precludes visiting $2- Conversely, if the agent deterministically takes action a = 1 at it will never visit thereafter.
Proposition 3.1. Hazan et al. 2019) There always exists stationary Markovian policy that max- imizes the entropy of state stationary distribution (i.e., solution of 31). Such an optimal stationary Markovian policy is generally stochastic.
Lemma C.1. (Puterman 1994) For any possibly non-Markovian policy n, define a stationary Markov policy n' as (a)s)) dz(s.a Then, dx dz
d=(s,a)"
Proof. Firstly, it is evident that there must exist least one non-stationary, non-Markovian pol- icy capable of maximizing the entropy of the state's stationary distribution. Then, according to Lemma C.I we can always obtain its corresponding stationary and Markov policy. This derived policy is generally stochastic, any deterministic policy can be arbitrarily bad in terms of maxi- mization of state entropy, as exemplified F
C.2 PROOF OF THEOREM 3.2
Theorem 3.2. The minimax optimization problem strained minimization problem:
p is equivalent solving the following uncon-
min Ep [u(s))] Egp [of4 (s Σ Exx(((mu6s 1) -[V. (10) k.
16

---

Published as a conference paper at ICLR 2025
where f4(y) maX+20 xy fzz). The objective function L(v, (1) is convex for and μ. Fur- thermore, given the optimal solution (v", ,p²) = arg min,, LVv, H.), the stationary distribution corrections of the optimal SEM policy given by:
d' (s. dD (8, a)
(1-)- (²)))) BB -(0))))))))))))) =: =: (s,
(11)
(9,0))
where 2,4 == max
Proof. Start with Equation 9
min max (1 Boo [u((u)] - al220 [
d(a.a)
)}
Σ d(s,, (),a,)) [(dss + log d(s))
(P
MR 8,400
=:w(a,n)
min max (1 - En [u[so]] a/E' [I(w(s,.))] + Σ d ds
d(s)
d(s) ((s(s())) + log d(s))
Pu d,w20
(23)
First, we can derive the flosed-form solution for d in 23)-
dL (u(s) log d((s) -
(24)
Od(s)
d" (s) exp(-u((s)-) 1)
(25)
Check the second order conditions to confirm that L is concave with respect to d(s):
LL 1 od(s)²
(26)
d(s)
We can plug this solution to L, which eliminates the max
min max(1 - E., [M(s)]] - aE [(w(s,a)) 1 w(ss ()), + Σ exp((=M(s - 1)
bip #20
(27)
Consider simplifying the second term in (27 using Fenchel conjugate:
max - 1.
(w((s, 1 w a)em
(28)
w20
max
aa A. w20
a
w(s aussaa)eu, a) - (----s(wi (w((s
(29)
alEgp
max wos a) () (d [(w(s,a)) mm(aa)20
(30)
-EA [f (s,e)((.
(31)
Hence, the Equation is equivalent to:
min (1 - Ep. + EaD [ (aff(ieeeae (s)))() + Eexp(-H(s) - -1) =: L(v, 41
Consider obtaining the closed-form solution for the minermmaximization for w 27
0L Ow(s,a) a) -add (s, a) [woss a)] + dP dl(s, (s,a)eenn
(32)
+ (8, = pr (2f(uu((())))) where := m)
(33)
17

---

Published as a conference paper at ICLR 2025
Again, plug w" (s, a) into the original objective function, which results in the minimization problem:
min - Bs0 [m(so)] - aB [s(- ((na(taa) 2i
( ),) - _ ( enu(s,
(Eua(saa) a
Σ exp(-M(s)
We can check that the function L is convex with respect to both and μ., by noting that f+() is convex function.
C.3 PROOF OF THEOREM 3.3
Theorem 3.3. Define the objective functions L(v, (1) as
T(y. p) == (1 r)) [(())] + F Es (oft (eeaa(,aa) + log E exp(-n(()) (12)
Then, for any optimal solutions (v" p²)) arg miny, [(v, H)) and (v", ,T*) arg min,, [V, p)) the following holds:
L(v",T'" TT*)
(13)
Also, there exists a constant such that the following holds:
μ¹" = p C. and v² V0' + -
(14)
Lemma C.2. For any functions v and μ, and constant C, the following equality holds:
T(,,u) = L(v + C S.
(34)
L(v + +C)
(35)
=11 110))s0-mll(t + [ + + log Σ exp((-pi - -C)
(36)
+ E ( af ((((())++ 3B + pen 2 - vvs) - -11
= (1 - 7) E Esg [((o))] + 1 E (sa) aas ( + +log] Σ exp((1u(s)) (1s)) - -l
(37)
(1 - Ee + Bs aff
(ate (e)) + log Σ Eexp(-M(s))
(38)
-V.
(39)
Lemma C.3. For any functions v and μ. L(v, μ)) 2 L(v, μ)) holds. The equality holds if and only if
Σ expp(-ule u(s) 1) =.
(40)
Proof. For any æ 2 0, the inequality x > log x always holds, where the equality holds if and
only i 1. By applying this equality, we have:
Σ exp(-u(s) - 1) - 2 log Σ exp(-M(s)) -1)
(41)
exp(-H() 1) - log Σ exp( (ap(-u(s) exp(-Ij
(42)
exp( (-((s) 1) > log Σ exp(-u(s))
(43)
L(V n) > L(v, 4)
(44)
where the equality holds if and only if expl expeppp(19
-p(s)
18

---

Published as a conference paper at ICLR 2025
Lemma C.A. Let (V" arg I(v, p.). Then, there exists constant C such that (v- + C, + 1-7) is an optimal solution of ming, L(v, jr).
Proof.
Let C = log exp(-"((s)-11 be constant. Also, let = = I- and p 11" CC*
(s)
=
Then,
Σ (xxp(x((s) - 1) = Σ exp
(x)
Σ exp (-77 1 log Σ exp( - -1)
Σ. expl --7' 1) = 1.
(45)
E of cxp( -1 - 1)
Then, have
L(² = L(F", x²)
(by 45 and Lemma C.3
(46)
L(w", )ir)
(by Lemma C.2
(47)
min [(v, 14) 2g
(48)
S min (v,, up
(by Lemma C.3
(49)
which implies that (8", FT)) is the optimal solution of min,,p L(v, p4)
D F-DIVERGENCE
In the experiments, we use the softened version of chi-square divergence for f:
x (x) = x log ++1 1
if x < if 1r21. 1.
(Jsont (x)') (z)) = (exp(x) if <
(x 1)2
+1 if 2 0
When f = Isolexx2 its corresponding fi(v) = max Ty (xx is given by:
8
i 0
(exp(y)
if < 0
f+ log ++
if < <
fi(1)
(50)
1202 +
yf y
-x2
if 2
where fI. is the Fenchel conjugate of f+-
E HYPERPARAMETERS
Baseline hyperparameters taken from URLB (Laskin et al. 2021 except for using the smaller hidden sizes 1000 + 256 for fast training/evaluation.
SEMDICE pretraining hyper-parameter
Value
Action repeat
a:
0.5
Batch size
1024
f-type
softchiq
Hidden dimension
256
Learning rate
0.0001
nsteps
Update every step
2
Table 1: Hyperparameter Settings
For the policy network, Gaal-Gaussinn distributions, following the baselines in URLBB
19

---

Published as a conference paper at ICLR 2025
F POLICY EXTRACTION
Our practical SEMDICE optimizes T6) that yields v² and μ². However, D² itself is not a directly executable policy, so we should extract a policy from them. Note that the optimal SEM policy x. is encoded in w":
d' ds = (1 1 (6) + BBr (x0) (s)
w (8)
=6p .p (s,a)
Then, we extract policy from w" via I-projection policy extraction method (Lee et al. 2021).
min IKL (d
(51)
[log w" (s. a) - DrL((([as)) +
(52)
log(f") cv (s). (2(r))))((( - -Duula(ao))oDllss +
= -DandD aro
(53)
E F
log(f")) (- +
(54)
amot
where C' is some constraint and (a|s)) is data policy. As perform online optimization, aD can be considered as an old policy, and we ignored the KL term in our practical implementation. Finally, enable ev- .u to be evaluated every action a, train additional parametric function (implemented as an MLP that takes (s, a) as an input and outputs a scalar value) by minimizing the mean squared error:
min (s [(((s,a) a) evr e (s. s')
(55)
In the following section, we present the pseudo-code for practical SEMDICE.
20

---

Published as a conference paper at ICLR 2025
G PSEUDO-CODE FOR SEMDICE
To sum up, SEMDICE computes a SEM policy by optimizing (v", which corresponds to ob- taining stationary distribution correction ratios of the optimal SEM policy. Then, we extract policy by training a e-network and performing I-projection as described in the previous section.
We assume V; 14, e, π parameterized by w, vi & respectively. Then, we optimize the parameters via stochastic gradient descent. The loss functions are summarized in the following:
[ == E) [vo()))] + ((s aff [dem (s, a, s,s
+ log E (d)
- log d0 (s))
(56)
J(e)) : (s,a, (ep(s,a) evo.u a, (s,
(by (55)
(57)
(TT) Bn log(f")- ((og(())((eee(ssa)))
(by 54
(58)
where evo.u (8, 8") (w) + 7w0( density estimation using k-nearnest-neighhor: - log dD (x:) R log (Il x
vo For log d (s) in 56 we used particle-based
Instead of
fully optimizing vo and Mou until convergence, we alternatively perform single gradient updates for
Mw, and Tg. The pseudocode of SEMDICE is presented in Algorithm
Algorithm SEMDICE
Input: Neural networks ve. Mus and By with parameters 0, and vh policy network wy with parameter Φ, replay buffer D, a learning rate 7), regularization hyperparameter a
1: for each timestep & do
2: Sample an action a, ~ To (5x) for the current state St-
3: Observe next state St+1 P(-|s;, ax) by taking action to the environment.
4: Add transition replay buffer D DU (st, ar 5t+1)
5: Sample minibatch from D for the following SGD updates.
6: Perform SGD updates:
(0,w) (0,w) Vo (vo, (w)
(Eg.
(56)
vb + th - 1/ ((w))
(Eg.
57
Φ Φ (TTO) (Eq.
58
7: end for
"Though practical also (approximate) estimation of do) via KNN, it does not contradict the main motivation of paper. (56). we use Monte Carlo integration approximate log
log Egp (exp(-u(s) log d (s))). While this approximation may introduce bias, still aims to maximize the entropy of the target policy's 's state stationary distribution from an arbitrary off-policy dataset, rather than maximizing the state entropy of the replay buffer.
In contrast, existing methods focus maximizing the intrinsic reward f(s) Rs log d (s). which cor- responds maximizing the state entropy of the replay buffer. That being said, they approximating log d' (s) by - log dz (s). As result, existing methods optimize the biased objective even when the kNN approximation error vanishes, whereas SEMDICE optimizes the unbiased objective once the kNN approxima- tion error vanishes. See also Appendix ! which shows that SEMDICE can compute the optimal SEM policy even with dataset collected by uniform random policy. This behavior of computing optimal SEM policy from random datasets cannot he achieved other baselines.
21

---

Published as a conference paper at ICLR 2025
H ADDITIONAL EXPERIMENTS ON TABULAR MDPS COMPARISON WITH MAXENT (HAZAN ET AL., 2019)
Uniform
Optimal
MAXENT
SEMDICE (ours)
Entropy of (s)
Entropy of (s)
1.0
normalized entropy
normalized sealzized am
0.5
0.0
0.0
10
10
10a
10
10
the number episodes
the number of episodes
Figure 7: Performance of SEMDICE and MaxEnt (Hazan et al., 2019) in randomly generated tabular MDPs. The first column indicates normalized policy entropy (0: state entropy of uniform policy, I: state entropy of optimal SEM policy), the second column indicates normalized entropy of the cumulative experiences.
We did not compare SEMDICE with MaxEnt (Hazan et al. 2019) in the main text because Max- Ent is not designed to compute single policy but rather to obtain a of policies. Specifically, it maximizes the entropy of the state distribution resulting from the weighted average of state dis- tributions visited by multiple policies. This approach does not align with the scenario we consider, which focuses on fine-tuning single policy derived from pre-trained policy. Still, for quantitative comparison with MaxEnt, we have conducted additional experiments in tabular MDPs. the result, MaxEnt (with tuned hyperparameters) underperforms SEMDICE in terms of learning efficiency. We also would like to emphasize that SEMDICE computes only a single stationary Markov policy, whereas MaxEnt continuously increases the number of policies it maintains, making it computation- ally inefficient.
I ADDITIONAL EXPERIMENTS ON TABULAR MDPs SEMDICE USING DATASET COLLECTED UNIFORM POLICY
Unifarm
Optimal
SEMDICE
Entropy of (*())
Entropy d?((s)
2.0
namalized entropy
nomallize lloed entropy
0.5
0.0
0.0
10
J0a
10
10
L0
10
101
the number of episodes
the number of episodes
Figure 8: Performance of SEMDICE when the dataset is always being collected by uniform random policy randomly generated tabular MDPs. The first column indicates normalized policy entropy (0: state entropy of uniform policy, 1: state entropy of optimal SEM policy), and the second column indicates normalized entropy of the cumulative experiences. This result demonstrates the capability of SEMDICE that can be optimized from arbitrary off-policy experiences.
22

---

Published as a conference paper at ICLR 2025
J ADDITIONAL EXPERIMENTS ON TABULAR MDPS: VALUE-BASED BASELINES
In Figure D we presented the performance of the policy-based baselines: for each iteration, they compute policy gradients for the intrinsic reward and perform policy gradient ascents. In this section, provide an additional result for value-based RL baselines. The intrinsic rewards are defined in the same way as described in the main text.
(59) where N(s, a) is the cumulative (s, -vvisitaatio counts by the agent, N(s) is the s-visitation counts by the agent, and d (s) is the empirical state distribution of the cumulative experiences. Then, for each value-based basclines, they maintain Q-table (initialized by zeros) and perform Q-learning updates by:
CB-SA f(s, = V CB-S : f(s,a) VN(3) PB-S : f(s, a) = - logd ds)))
((s a) + Q(s.a) n (f(..)) + y ymx maxa(ssssseea d)) Q(s,a))
(60)
a
where 7, E (0, 1) is a learning rate.
As value-based RL methods do not maintain explicit policy, we evaluated two types of target pol- icy induced by Q: greedy (deterministic) policy (a(s) = arg max, Q(s, a)) and softmax policy (T(aas) x exp(Q(ss, a)/ For exploration, all methods adopt an e-soft behavior policy with e x (1/TT) decreasing over time. We also decrease the temperature softmax policy over time with the rate of (I/TT). The result for greedy policy is shown in Figure and the result for the softmax policy is presented in Figure 10
Uniform
Optimal
SEMDICE (ours)
CB-SA (greedy)
CB-S (greedy
+ -B-S (greedy)
Entropy of (s)))
Entropy of (s))
G G- -e----e----- - 0
nomaaiia antropy
0.0
normalized entropp
-0.0
20
1t
10²
101
111
10a
10
211
the number of episodes
the number of episodes
Figure 9: Performance of SEMDICE and value-based baselines in randomly generated tabular MDPs, where the target policy (s) of baseline is given by the greedy policy w.r.t. Q(s, a).
Uniform
Optimal
SEMDICE (ours)
CB-SA (saftmax)
CB-S (softmax)
+ PB-S (softmax)
Entropy of (s)
Entropy of (s)
soreallzeed entropy
normalics lizel entroyy
207
TiI
the number of episodes
the number of episodes
Figure 10: Performance of SEMDICE and based baselines in randomly generated tabular MDPs, where the target policy (als) of baseline is given by the softmax policy w.r.t. a).
As can be seen from the left column of Figure g the greedy target policy for the estimated Q exhibits very low state entropy. This result is expected, as the optimal SEM policy should generally be stochastic (Proposition 3.I. but the greedy target policy deterministic. Thus, it does NOT converge to the optimal stochastic SEM policy, leading significant suboptimality. The entropy of the cumulative state-visit experiences H[dd(s) (the right column of Figure 9) for PB-S approaches to the maximum state entropy, but it was achieved by mixture of many -tatioation determinstic policies. In contarst, SEMDICE yields a single stationary stochastic SEM policy.
In Figure 10 entropy for the target policy (left column). However, there is no guarantee that any of these value- based baselines with the softmax policy will converge to optimal SEM policy.
see that adopting the softmax (thus stochastic) policy leads to better state
23

---

Published as a conference paper at ICLR 2025
K
ABLATION EXPERIMENT RESULTS FOR SEMDICE
0.1
10.0
Entropy (s)s
Entropy of (s)))
1.0
normalized entropy
nomalized entropy
0.75
0.0
0.50
0.0
10:
the number of episodes
the number of episodes
Figure 11: Experimental results on varying (SEMDICE without function approximation).
Sx (a 0.1)
Suittxx (a 0.1)
JKL. {a 0.1)
Entropy of (F((s)
Entropy of (s)
1.0
normalized entropy
1.0
normalized mazaliddd entropy
0.9
0.8
0.4
10'
10a
10
10
10"
10
the number of episodes
the number of episodes
Figure 12: Experimental results on different f f-divergence (SEMDICE without function approxima- tion).
fe
state entropy (a
6.5
5.0
0.5
0.001
0.1
4
0,001
02
3.0
0.01
---
0.01
timesteps
timesteps
Figure 13: The effect of the varying a and on ContinuousMoumtaanCar (SEMDICE with neural function approximation).
As shown in Figure IE SEMDICE can become numerically unstable with very small values of a, but it remains not too sensitive to a as long as is within a reasonable range. Additionally, using different ff didn't 't make significant difference.
24

---

Published as a conference paper at ICLR 2025
L GENERALIZATION TO UNDISCOUNTED (7 = 1) SETTING
SEMDICE can be generalized to deal with
setting, where (s)
lim TमE Ei= Pr(st 8; =). To this end, the original optimization problem (3-5) in the
paper should be modified by adding an additional normalization constraint Σ., d(s, a) 1. Then, for any 7 E (0, 1),
max Lo(s)logd(s)) log)(s) aD, (d(s, aa)ld"
(61)
d,d20
s.t. _d ((s',a') = y)poo'(s + 2 d(s, () As
(62)
5,0
Σ d(x,a) = d(s)
Vs
(63)
d(s,,))=1 =1
(64)
By taking the similar derivation steps of SEMDICE with the additional constraint 64 show that solving the following unconstrained convex optimization problem is equivalent to solving 61 64). which additionally includes single scalar variable A:
we can
min (1 Eo (((s) EaD P.R.A.
[oft (2 emma)ss))) + + exp(-u(s) -
(65)
where x a) == A (u(s) + [x - vss) and À is Lagrange multiplier for the normal-
ization constraint (64). The detailed derivation will be included in the final version of the paper.
M EXTENSION TO UNDISCOUNTED (NON-STATIONARY) FINITE-HORIZON SETTING
SEMDICE be extended to a finite horizon setting, where the goal is to maximize the entropy of the average state distribution d= (s) == 1 El Pr(s; = s; x). The following formulation results
non-stationary (timeste-dependdent) policy by T. E. dj
(s,a)
max - E. de)) logd(s)) -a aTo D, (dels, a))lP(ss)) 4 4,420
(66)
s.t. Σ. do(s,, a) = (s())
(67)
vs
Σ. di(',,a''')) - E.a (d- ()T(sss))
(68)
Σ. () de(s) Vs,t E ().
(69)
We show that solving (66(69) is equivalent solving the following unconstrained convex optimization problem, where the difference to the original objective function is that v, l are timestep- dependent (i.e. Vt. Pc).
min (1 (1-7)te pEp [s0) Et Eao (afi (He EI E.exp exp(-Hi(s) 1) (oo (7
(70)
where (s a) == (s)+))+++)+))))+)))) EE. ve and vr+1(1) (.) = 0. The detailed derivation will be included in the final version of the paper.
25

---

Published as a conference paper at ICLR 2025
N MORE DISCUSSIONS ON SEMDICE'S OBJECTIVE
Despite its inclusion of f-divergence regularization, SEMDICE is not fundamentally opposed to particle-based state-entropy estimation; rather, addresses the same objective of maximizing state entropy, with the additional inclusion of f-divergence regularization to enhance the stability of the learning process. It is important to note that this regularizer not intended to impose strong constraint on the optimal policy. and its influence can be controlled by adjusting the hyperparam- eter a. Note that many successful standard reward- maximizing RL algorithms have made use of regularizations such entropy (Haarnoja et al. 2018) and KI.-divergence (Schulman et al.) 2015), f-divergence (Nachum et al. 2019b), greemmnn-divergence (Geist et al.) 2019), and so on. Although these regularizations might appear contrary to pure reward maximization, their inclusion actually stabilizes the overall learning process, and improves the overall reward performance of the resulting policy. In similar vein, state-entropy maximization method employs f-divergence regulariza- tion to prevent abrupt distribution shift (analogous to TRPO), enhancing the robustness and stability of learning process. Additionally, from technical perspective, f-divergence regularization makes the objective function strictly concave, guaranteeing the global optimality of any local optimum (i.e. ensuring the uniqueness of the optimal solution).
Also, SEMDICE's objective function indeed aims to maximize the state entropy, i.c. encourag- ing exploration of unvisited states more. To see this more intuitively, consider the main objective function of OptiDICE (Lee et al. 202I), a reward-maximizing RL algorithm:
[as= (+ (1 + +w(()) - -(s))]
min(1 - 3E (+0)] + )
(71)
Then, the following is the main objective function of our SEMDICE, state-entropy maximization method:
min(1 - ( E. [(())]] ((e (af= (a=())((s)) (#((() + m(65 - -(6))=(ss)) + log Σ exp(-M(s))
(72)
That being said, the for fixed μ. SEMDICE can be interpreted as OptiDICE with the reward function defined by pls). Of course, we are jointly optimizing the μ instead of using fixed μ. Also, due to the second term of the objective function (f is increasing function), p(s) is pressured to decrease more in regions of high data density and less where the data density is low. Conse- quently, SEMDICE's resulting policy will tend to explore unvisited states more actively, which is well-aligned with the goal of state entropy maximization.
O MACHINE AND SETUP
We run the experiments on machines with Titan XP GPUs. Pretraining took approximately 10 hrs and finetuning took around 30 minutes.
26


Reviews:
[{"review_id": "ZMSRn2PUdF", "reviewer": "Reviewer_R8AG", "summary": "The work proposed a novel state entropy maximization method called SEMDICE for RL pre-training. SEMDICE addresses the limitations of existing methods by directly optimizing the stationary distribution space rather than the policy or Q-function space, utilizing arbitrary offline samples. The study demonstrates that SEMDICE converges to the optimal SEM policy in tabular MDP experiments and is more efficient in RL policy pre-training compared to existing data-driven (i.e., SEM-based) unsupervised RL methods.", "questions": "* Including the algorithm flowchart within the main text would enhance readers' understanding of the algorithm.\n* Is the algorithm sufficiently robust to the hyperparameters, especially the $\\alpha$? Is there any method can be employed to mitigate adverse effects?\n* How is the computational efficiency in high-dimensional data?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The proposed algorithm SEMDICE is a novel method suitable for off-policy training.\n* SEMDICE demonstrates efficient convergence in tabular MDP experiments and outperforms existing methods in RL policy pre-training.\n* Theoretical analysis is provided to demonstrate the efficacy of the SEMDICE.", "weaknesses": "* The hyperparameter $\\alpha$ may have a significant impact on the results.\n* Though experiments in tabular RL and continuous state spaces are conducted, it does not address high-dimensional data (such as image-based tasks), which may pose computational complexity issues. This could be a potential direction for future improvements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZXAyUgGHsM", "reviewer": "Reviewer_6phi", "summary": "This paper focuses on developing an effective pretraining policy in an unsupervised reinforcement learning setting by maximizing state entropy. Previous methods either utilized k-nearest neighbors (kNN)-based particle entropy estimation, which is biased, or were based on the on-policy algorithms from Hazan et al. (2019), which are sampling inefficient. The authors build on the DICE (Stationary Distribution Correction Estimation) series of methods to propose a new off-policy entropy maximization approach named SEMDICE. The paper demonstrates the advantages of their method on benchmarks including Gridworld, MountainCar, and URL Benchmark.", "questions": "1. This paper addresses Convex MDP problems (where state entropy is convex with state visitation distribution). Is there potential for applying the DICE method to other Convex MDP problems, such as diverse skill discovery, as summarized in Table 1 of [2]?\n2. The experimental results indicate that the CIC method performs poorly in state coverage during the pretraining phase but excels in downstream task fine-tuning. This suggests that maximizing state entropy may not be the optimal pretraining strategy. Could you provide clarification on this point?\n\nReference：\n\n[2]  Reward is Enough for Convex MDPs. NeurIPS 2021.\n\n--- \nSome concerns have been addressed, I have raised my socre.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The motivation is well. The particle entropy method operates off-policy but optimizes the entropy of the replay buffer rather than the target policy. In contrast, the method proposed by Hazan et al. (2019) is on-policy but sample inefficient. Using DICE-like approaches can address these issues.\n\nWhile the original DICE framework focuses on Linear MDPs (where the objective function is linear with respect to the state visitation distribution), this paper examines Convex MDPs (where state entropy is convex with respect to the state visitation distribution). This shift introduces differences from the original DICE algorithm, such as the emergence of an exponential term. To mitigate gradient explosion, the paper employs a log-sum-exponential replacement for this term and theoretically demonstrates that this substitution does not significantly impact optimality.\n\nExperimental results indicate that the proposed method, which emphasizes state coverage, outperforms other methods, thereby validating its effectiveness.", "weaknesses": "1. The work of Hazan et al. (2019) is a primary reference, yet the paper lacks a comparative analysis of its experiments, particularly like the MountainCar experiment.\n2. The implementation details are vague. For instance in Equation 56 of Appendix, the paper estimates $−log d^D(s)$ using a k-nearest neighbors (kNN) based particle entropy estimation, which introduces bias and contradicts the paper's motivation. This needs clarification and should be discussed in the main text.\n3. For readers unfamiliar with DICE, the paper is quite opaque. I strongly recommend that the authors organize the theoretical results and the main derivation process similarly to Appendix B in [1], rather than presenting a cumbersome Theorem (e.g., Theorem 3.2).\n\n**Other Typos:**\n- Line 49: Missing space before \"However.\"\n- Line 296: \"inimizing\" should be corrected to \"minimizing.\"\n- Line 298: The numbering (1), (2) is likely to cause confusion.\n- Appendix A: The meaning of $d_+$ is not explained.\n\nReference：\n[1] Dual RL: Unification and New Methods for Reinforcement and Imitation Learning. ICLR 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kdi4o8L3jz", "reviewer": "Reviewer_9ykY", "summary": "This paper tackles the challenge of unsupervised pre-training in reinforcement learning, aiming to develop a versatile policy that can generalize to downstream tasks without relying on specific reward signals. The authors focus on maximizing state entropy (SEM) to derive a policy that promotes diversity in the agent's experience by increasing the entropy of its stationary state distribution. They introduce **SEMDICE**, a structured off-policy algorithm that derives an SEM policy from arbitrary off-policy data by directly optimizing within the space of stationary distributions. SEMDICE produces a single, stationary policy that effectively maximizes state entropy. Experimental results indicate that SEMDICE surpasses existing baselines in enhancing state entropy and demonstrates superior efficiency in adapting to downstream tasks, making it a strong candidate for SEM-based unsupervised RL pre-training approaches.", "questions": "Similar to the issue of weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper addresses the limitation of existing SEM methods in their inability to support off-policy learning, introducing the objective of optimizing the state distribution. By integrating optimization techniques from the DICE family of methods, the authors propose a well-motivated and logically consistent algorithm.  \n2. The mathematical derivations in this paper are both accurate and well-founded.  Specifically, this paper provides a detailed explanation of why it is necessary to introduce both $d(s,a)$ and  $\\bar{d}(s)$, which is highly convincing from an algorithm design perspective.\n3. The experiments demonstrate that the SEMDICE method outperforms all baselines.  \n\nIn summary, this paper presents a clearly motivated approach, complete algorithmic derivation, and comprehensive experimental analysis.", "weaknesses": "1. I think there may be an oversight in the derivation. Specifically, in Equation (15), the summation over $s$ is replaced by sampling $s$ from $D$ through importance sampling, which in fact overlooks cases where $s$ is not in $D$. Therefore,  if $s$ is severely out-of-distribution (OOD), the calculated $ \\log \\bar{d}(s) $ may be highly inaccurate. Consequently, even though SEMDICE is off-policy, it still introduces a significant bias. The authors might consider adding experiments to demonstrate the impact of OOD-induced bias on the algorithm's performance.\n\n2. In Section 5.1, the different baselines represent various intrinsic reward construction methods, all of which are based on policy-gradient algorithms. Why not use value-based methods as the core algorithms instead? This could help avoid the gradient bias introduced by off-policy policy gradients. Since SEMDICE’s off-policy nature does not suffer from gradient bias, using value-based methods would allow for a fairer comparison between SEMDICE and other algorithms.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VAnBS107Vx", "reviewer": "Reviewer_qDon", "summary": "This paper addresses the issue of state entropy maximization by off-policy and principled method. \nKey idea underlying this work is the Eq.4-6, where the authors jointly optimize state stationary distribution and state-action stationary distribution.\nThe method is evaluated in both toy examples and tasks in URLB.", "questions": "1, As described in Eq.57, the performance of network $e_{\\phi}$ is heavily dependent on the optimality of network v and u. This design may lead to unstable training.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1, There is a theoretical analysis (Lagrange dual formulation and convex problem simplification) for the proposed approach. While I have not checked the proofs carefully, they appear to be correct.\n\n2, The goal of the paper for addressing the challenge of state entropy maximization in RL is well motivated.\n\n3, The results in the experiments clearly support the claims made in the paper.", "weaknesses": "1, Why the performance (entropy index) of APT and CIC drop after training in Fig.3 in Quadruped experiment?\n\n2, Will the code be publicly available after acceptance? Some details of the specific implementation are not given.\n\n3, Line 296: inimizing -> minimizing", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)