强化学习小笔记 —— 从 Normal 正态分布的对数概率密度到 tanh-Normal的对数概率密度

在学习 SAC 算法用于连续动作的代码时，遇到了一个不懂的地方，如下代码所示：

# pytorch
class PolicyNetContinuous(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound):
        super(PolicyNetContinuous, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc_mu = torch.nn.Linear(hidden_dim, action_dim)
        self.fc_std = torch.nn.Linear(hidden_dim, action_dim)
        self.action_bound = action_bound

    def forward(self, x):
        x = F.relu(self.fc1(x))
        mu = self.fc_mu(x)
        std = F.softplus(self.fc_std(x))
        dist = Normal(mu, std)
        normal_sample = dist.rsample()  # rsample()是重参数化采样
        log_prob = dist.log_prob(normal_sample)
        action = torch.tanh(normal_sample)
        # 计算tanh_normal分布的对数概率密度
        # 我们需要的是进行 tanh 转换后的对数概率密度，不是 normal 对数概率密度
        log_prob = log_prob - torch.log(1 - torch.tanh(action).pow(2) + 1e-7)
        action = action * self.action_bound
        return action, log_prob

# tensorflow
class PolicyNetContinuous(tf.keras.Model):
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound):
        super(PolicyNetContinuous, self).__init__()
        self.fc1 = tf.keras.layers.Dense(hidden_dim, activation=tf.keras.activations.relu)
        self.fc_mu = tf.keras.layers.Dense(action_dim)
        self.fc_std = tf.keras.layers.Dense(action_dim, activation=tf.keras.activations.softplus)
        self.action_bound = action_bound

    def call(self, x):
        x = self.fc1(x)
        mu = self.fc_mu(x)
        std = self.fc_std(x)
        dist = tfp.distributions.Normal(loc=mu, scale=std)
        normal_sample = dist.rsample()
        log_prob = dist.log_prob(normal_sample)
        action = tf.tanh(normal_sample)
        log_prob = log_prob - tf.math.log(1 - tf.tanh(action) ** 2 + 1e-7)
        action = action * self.action_bound
        return action, log_prob

疑惑的地方在于对动作的对数概率密度进行了以下处理：

log_prob = log_prob - torch.log(1 - torch.tanh(action).pow(2) + 1e-7)

从代码里可以看到，策略网络的目标是输出一个动作，这个动作是从一个参数为 mu 和 std 的正态分布中采样得到的。然后，这个动作被一个 tanh 函数转换到 -1 到 1 的范围内，以满足环境的动作空间要求。

这里的 log_prob 是动作的对数概率密度。在 SAC 算法中，我们需要计算这个对数概率密度，因为它在更新策略网络时会用到。

因为动作是从正态分布中采样得到的，然后使用tanh()进行转换，经过转换后的动作的对数概率密度不能再使用从正态分布中获取的log_prob。所以，我们需要计算的是这个转换后的动作的对数概率密度，而不是转换前的动作的对数概率密度。

下面来说明一下，基于转换后的对数概率密度是怎么得到的：
（1）首先，我们从一个正态分布$Normal(mu,std)$中采样一个随机变量$x$，其概率密度函数为$p_X(x)$，接着，我们使用$tanh$函数将随机变量$x$映射为另一个随机变量$y=tanh(x)$，y对应的概率密度函数为$p_Y(y)$。
（2）根据反函数的定义，我们可以知道有：$y=tanh(x)$ 和 $x=tan{h}^{-1}(y)$，其中$tan{h}^{-1}为tanh$的反函数。
（3）我们考虑$p_Y(y)$对应的累积分布函数：$$F_Y(y)=p(Y<=y)=p(g(X)<=y)$$ 公式中的$g=tanh(\cdot )$；
（4）根据反函数的性质和tanh(·)的单调性，可以将(3)中的公式修改为：$$F_Y(y)=p(X<=tan{h}^{-1}(y))=F_X(tan{h}^{-1}(y))$$。
（5）要得到$Y$的概率密度函数$p_Y(y)$，我们可以对$F_Y(y)$进行求导，根据链式法则有：$$p_Y(y)=\frac{d{F}_X(tan{h}^{-1}(y))}{dy}=\frac{d{F}_X(tan{h}^{-1}(y))}{dtan{h}^{-1}(y)}\cdot \frac{dtan{h}^{-1}(y)}{dy}=\frac{p_X(x)}{1-(tanh(x){)}^2}$$
（6）获取对数概率密度函数：$$log{P}_Y(y)=log{P}_X(x)-log(1-(tanh(x){)}^2)$$
（7）为了避免$1-(tanh(x){)}^2$等于0导致log趋向于无穷小，所以一般会加一个$1e-7$，即：$$log{P}_Y(y)=log{P}_X(x)-log(1-(tanh(x){)}^2+1e-7)$$
（8）综上，得到了代码中的结果。