which is not straightforward in practice.

To address this, we estimate the state distribution by separately training a density estimator based on sampled states in the rollout. The crucial step here is that, we use a density estimator that is explicitly a function of the policy parameters $\theta$. In other words, our density estimator takes as input, the parameters $\theta$ of the policy itself (for instance, weights of a policy neural network) through which we now obtain an estimate of $p_{\theta}: \theta \mapsto \Delta(S)$, where $p_\theta(s)$ is the occupancy probability (discounted or otherwise) of state $s$. We use a variational inference based density estimator, which can be trained to maximize a variational lower bound to the the log-likelihood of $p_{\theta}(s)$. As a result, we can obtain an estimation of $d_{\pi_{\theta}}$ since in case of stationary distributions, we have $\log p_{\theta}(s) = \log d_{\pi_{\theta}}(s)$. Estimation of $d_{\pi_{\theta}}$ under any policy $\pi_{\theta}$ requires collecting a large number of samples from the rollout. Instead of this, we can use ideas from multi-scale stochastic algorithms to learn this in an online manner. Hence, we require a separate time-scale for training the density estimator, in addition to learning the policy and value functions in policy 