How to build the actor policy of Soft-Actor-Critic after sampling from a Multivariate normal distribution?

Question

I'm trying to solve LunarLanderContinuous-v2 (https://www.gymlibrary.ml/environments/box2d/lunar_lander/) using Soft Actor-Critic algorithm (following the pseudocode above)

To update the actor policy (step 14) I understood that I need to:

1. sample an action from a normal distribution using mean and variance that are the output of the policy network
2. squash the action with tanh to have action bounds [-1, 1]
3. compute the log_prob, the log probability density of the sample
4. do the following operation in tensorflow

# STEP 14
bounded_actions_sample, predicted_next_actions_sample = policy_network.predict_actions(batch_observation)
pdf = get_pdf(predicted_next_actions_sample, batch_observation)
q_min = tf.math.minimum(q_network_1(batch_observation, bounded_actions_sample), q_network_2(batch_observation, bounded_actions_sample))
        
loss_policy_network = tf.reduce_mean(alfa * tf.math.log(pdf) - q_min))

where "pdf" is the value of the probability density function of the normal distribution computed from predicted action (not clipped), mean and variance.

Since the action space of LunarLanderContinuous-v2 is continuous and the possible actions are 2, my idea is to predict means μ, variances σ and a correlation ρ parameter in order to build a Multivariate normal distribution using tfp.distributions.MultivariateNormalTriL and sample actions from it.

My problem is that, according to the Appendix C of the paper https://arxiv.org/pdf/1801.01290.pdf, I have to compute the likelihoods of the bounded actions using

but I don't know how to do that because the first term is a single number due to it come from a Multivariate normal distribution, instead the second term contains 2 elements, one for each action. Has anyone any idea to code it?

Answer 1

I don't really understand the confusion, as both terms are scalars (where $u_i$ are the components of $u$).

So if $u\sim \mathcal{N}(\mu, \Sigma)$, then the first term would be

$$\log \mu(u\mid s)=-\frac{D}{2}\log(2\pi)-\frac{1}{2}\log\vert \Sigma\vert - \frac{1}{2}(u-\mu)^T\Sigma^{-1}(u-\mu) $$

However, SAC assumes that $\Sigma$ is diagonal, which makes the expression easier, that is

$$\log \pi(a\mid s)=-\frac{D}{2}\log(2\pi)-\frac{1}{2}\sum_{i=1}^D\log\sigma_i - \frac{1}{2}\sum_{i=1}^D\frac{(u_i-\mu_i)^2}{\sigma_i}-\sum_{i=1}^D\log(1-\tanh ^2u_i) $$

If there is a correlation between the individual actions, you would have to implement the first expression.

An explicit implementation I have come across in Pytorch is the following: (where $\sqrt {2\pi} \approx2.5066$)

sampled_u = self.sample_z(mua_tensor.data, siga_tensor.data).data
sampled_a = torch.tanh(sampled_u)

log_pi_exp = torch.sum(- (mua_tensor.data - sampled_u.data).pow(2)
                       / (siga_tensor.data.pow(2)) / 2
                       - torch.log(siga_tensor.data * torch.tensor(2.5066)),
                       dim=-1, keepdim=True)
log_pi_exp -= torch.sum(torch.log(1.0 - sampled_a.pow(2) + EPS), dim=-1, keepdim=True)

However, this is not efficient and there are undoubtedly specific functions in Tensorflow that make this implicitly, but I hope this provides more intuition.