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

We thank the reviewer for the thoughtful feedback.

[Weaknesses]

1. When we visualized the learned policies of APT, we observed that it remained in an upside-down position while lying on the ground after the pre-training process of the Quadruped. This behavior results in a small state entropy. In contrast, other methods including SEMDICE exhibit much more diverse and dynamic behaviors, such as jumping, getting up, and lying down.

2. Yes, the code will be publicly available.

[Question]

1. As the reviewer mentioned, $e_\phi$ should ideally be optimized using the converged $\nu$ and $\mu$. However, in the context of deep RL, it is common practice to perform alternating gradient updates between two networks rather than waiting for a dependent network to converge. For example, actor-critic algorithms like TD3 alternate between gradient updates for the actor (using the inaccurate Q) and the critic networks, even though the actor update should ideally be based on an accurate $Q^\pi$ estimation. It can be understood that SEMDICE follows a similar approach, and it has been shown to work well in practice.

We thank the reviewer for the thoughtful feedback.

[Weaknesses 1]

As the reviewer pointed out, SEMDICE only uses (s, a, s') samples from the dataset D for optimization. This procedure can be viewed as computing an optimal SEM policy in the maximum-likelihood estimation (MLE) MDP for D. However, this approach is not exclusive to SEMDICE; most off-policy RL algorithms can be understood as the process of computing an optimal policy in an MLE MDP. While it is true that this process is biased, SEMDICE's (in-sample) learning approach encourages the agent to visit more states that it has not previously encountered (see Appendix M for more discussion). As a result, the agent gradually expands the state space it explores and performs well in practice.

For the Monte Carlo integration (line 265), we also tested q(s) = Unif(s), which samples uniformly within the min-max range of each state dimension. However, this caused numerical instability due to the exploding values of $\mu(s)$. Specifically, in Eq. (15), the second term plays a role of reducing $\mu(s)$, while the third term increases it. However, since there is no learning signal to reduce the increased $\mu(s)$ values for OOD states sampled from $s \sim Unif(s)$, $\mu(s)$ easily grows to infinity especially when the state dimensions are large (like Jaco and Quadruped).

[Weaknesses 2]

Following the reviewer's suggestion, we have conducted additional experiments with the value-based baselines in tabular MDPs (see Appendix I in the updated paper). However, please note that the value-based baselines maintain a Q-table rather than an explicit policy, and the target policy is implicitly represented by the Q-values (e.g., greedy for Q or softmax for Q). We have provided both sets of results, but the value-based baselines still significantly underperform SEMDICE. This is expected for the following reasons: (1) For a greedy (deterministic) policy, it is suboptimal because the optimal SEM policy is generally stochastic (see Figure 6 and Proposition 3.1), and (2) for a softmax (stochastic) policy, there is no guarantee that such a naively randomized policy can achieve state entropy maximization.

We thank the reviewer for the thoughtful feedback.

[Weaknesses 1: Comparison with MaxEnt]

The primary reason we did not compare SEMDICE with MaxEnt (Hazan et al.) in our experiments is that MaxEnt is not designed to compute a single policy but rather to obtain a set of policies. Specifically, it maximizes the entropy of the state distribution resulting from the weighted average of state distributions visited by multiple policies. This approach does not align with the scenario we consider, which focuses on fine-tuning a single policy derived from a pre-trained policy. Still, for quantitative comparison with MaxEnt, we have conducted additional experiments in tabular MDPs [1]. In 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 computationally inefficient.

[1] https://i.ibb.co/J55vX1j/result-tabular-mdp.png

[Weaknesses 2: Particle-based entropy estimation and paper's motivation]

We would like to clarify that the (approximate) estimation of $-\log d^D(s)$ via kNN does not contradict the motivation of the paper. In Eq. (56), we use Monte Carlo integration to approximate $\log \sum \exp (-\mu(s))$ as $\log \hat{\mathbb{E}}_{d^D} [ \exp( -\mu(s) - \log d^D(s) ) ]$ (see line 265). While this approximation may introduce bias, it still aims to maximize the entropy of the target policy'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 on maximizing the intrinsic reward $\hat r(s) \approx - \log d^D(s)$, which corresponds to maximizing the state entropy of the replay buffer. That being said, they are approximating $-\log d^\pi(s)$ by $- \log d^D(s)$. As a result, existing methods optimize the biased objective even when the kNN approximation error vanishes, whereas SEMDICE optimizes the unbiased objective once the kNN approximation error vanishes. We will add this clarification to the paper. Please also refer to Appendix H, which shows that SEMDICE can compute the optimal SEM policy even with a dataset collected by a uniform random policy. This behavior of computing an optimal SEM policy from random datasets cannot be achieved by other baselines.

[Weaknesses 3: Paper presentation]

Thank you for your feedback. We are currently adding more detailed and organized theoretical results and derivations in the Appendix, referring to the mentioned paper, for readers' better understanding. These additions will be reflected in the final version of the paper.

[W: Typos]

Thanks for pointing out the typos, which are modified in the updated paper.

[Questions]

1. Yes, we believe that a skill-discovery algorithm (e.g., by maximizing mutual information) can be derived based on the DICE formulation. This would enable MI maximization between tasks and target policy's states using an arbitrary off-policy dataset, whereas existing off-policy methods (e.g. [2]) optimize the mutual information between tasks and replay buffer's states. Deriving a DICE-based method for skill-discovery is one of our future directions.

2. We would like to clarify that the optimal pretraining strategy may vary depending on the distribution of downstream tasks. In the case of URLB, only four predefined downstream tasks, designed by humans, are provided for each domain. State entropy maximization approach like SEMDICE would offer a robust initialization that is resilient to worst-case scenarios across a broader range of possible downstream tasks (see Appendix A for more discussion). However, this does not necessarily mean that state entropy maximization is the optimal pretraining strategy for all specific downstream tasks, especially when only a small number of tasks are considered as in URLB. Although CIC's state visitation coverage during pretraining was limited compared to SEMDICE, it can still perform well on certain downstream tasks. To sum up, while state entropy maximization provides a solid foundation for robust performance across diverse tasks, it is not always guaranteed to be the best approach for every specific downstream task.

[2] Liu and Abbeel, APS: Active Pretraining with Successor Features, 2021

We thank the reviewer for the thoughtful feedback.

[Robustnesss to the hyperparameters]

In Appendix J, we have presented ablation experiments to see the sensitivity of the choice of $\alpha$ and $f$. In short, using too small $\alpha$ incurs numerical instability (due to $\frac{1}{\alpha}$ term in the objective functions), while using too large $\alpha$ can slow down the training efficiency due to its conservative updates. Overall, as can be seen in Figure 10-12, with a moderately large range of the choice of $\alpha \ge 0.1$, SEMDICE exhibits robust performance.

[High-dimensional domains]

In the paper, we present experimental results on state-based URLB, where the state dimensions are 24 for Walker, 55 for Jaco, and 78 for Quadruped. They provide moderately high-dimensional continuous state space domains, where SEMDICE performs the best in terms of state entropy maximization. Please note that CIC [1] also conducted experiments on state-based URLB only.

Extension of SEMDICE to pixel-based domains remains a future work. Still, although the extension may require combining with an efficient representation learning technique (e.g. [2,3]), SEMDICE would not pose a significant challenge to computational efficiency in high-dimensional domains: SEMDICE optimizes state-dependent function parameterized by two neural networks $\nu_\theta: S \rightarrow \mathbb{R}$ and $\mu_\omega: S \rightarrow \mathbb{R}$ (vs. other methods learn $Q: S \times A \rightarrow \mathbb{R}$ function) with an additional image encoder $E_{\psi}$. Assuming the same encoder network is used as in other methods, the computational cost in training is expected to be similar to that of other methods.

[Algorithm Flow]

Thanks for your suggestion. We will include an additional simple flowchart to clarify the process. The overall flow of the practical algorithm implementation, along with the pseudo-code, can be found in Appendix G. Optimizing Eq. (58) corresponds to solving the original state entropy maximization problem defined in Eq. (4-6), while Eq. (59-60) describes the policy extraction from the optimized stationary distribution corrections.

[1] Laskin et al., CIC: Contrastive Intrinsic Control for Unsupervised Skill Discovery, 2022

[2] Fujimoto et al., A Deep Reinforcement Learning Approach to Marginalized Importance Sampling with the Successor Representation, 2021

[3] Yarats et al., Reinforcement Learning with Prototypical Representations, 2022