在学习 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)首先,我们从一个正态分布
N
o
r
m
a
l
(
m
u
,
s
t
d
)
Normal(mu,std)
Normal(mu,std)中采样一个随机变量
x
x
x,其概率密度函数为
p
X
(
x
)
p_X(x)
pX(x),接着,我们使用
t
a
n
h
tanh
tanh函数将随机变量
x
x
x映射为另一个随机变量
y
=
t
a
n
h
(
x
)
y=tanh(x)
y=tanh(x),y对应的概率密度函数为
p
Y
(
y
)
p_Y(y)
pY(y)。
(2)根据反函数的定义,我们可以知道有:
y
=
t
a
n
h
(
x
)
y=tanh(x)
y=tanh(x) 和
x
=
t
a
n
h
−
1
(
y
)
x=tanh^{-1}(y)
x=tanh−1(y),其中
t
a
n
h
−
1
为
t
a
n
h
tanh^{-1}为tanh
tanh−1为tanh的反函数。
(3)我们考虑
p
Y
(
y
)
p_Y(y)
pY(y)对应的累积分布函数:
F
Y
(
y
)
=
p
(
Y
<
=
y
)
=
p
(
g
(
X
)
<
=
y
)
F_Y(y)=p(Y<=y)=p(g(X)<=y)
FY(y)=p(Y<=y)=p(g(X)<=y) 公式中的
g
=
t
a
n
h
(
⋅
)
g=tanh(·)
g=tanh(⋅);
(4)根据反函数的性质和tanh(·)的单调性,可以将(3)中的公式修改为:
F
Y
(
y
)
=
p
(
X
<
=
t
a
n
h
−
1
(
y
)
)
=
F
X
(
t
a
n
h
−
1
(
y
)
)
F_Y(y)=p(X<=tanh^{-1}(y))=F_{X}(tanh^{-1}(y))
FY(y)=p(X<=tanh−1(y))=FX(tanh−1(y))。
(5)要得到
Y
Y
Y的概率密度函数
p
Y
(
y
)
p_Y(y)
pY(y),我们可以对
F
Y
(
y
)
F_Y(y)
FY(y)进行求导,根据链式法则有:
p
Y
(
y
)
=
d
F
X
(
t
a
n
h
−
1
(
y
)
)
d
y
=
d
F
X
(
t
a
n
h
−
1
(
y
)
)
d
t
a
n
h
−
1
(
y
)
⋅
d
t
a
n
h
−
1
(
y
)
d
y
=
p
X
(
x
)
1
−
(
t
a
n
h
(
x
)
)
2
p_Y(y)=\frac{dF_X(tanh^{-1}(y))}{dy}=\frac{dF_X(tanh^{-1}(y))}{dtanh^{-1}(y)}·\frac{dtanh^{-1}(y)}{dy}=\frac{p_X(x)}{1-(tanh(x))^2}
pY(y)=dydFX(tanh−1(y))=dtanh−1(y)dFX(tanh−1(y))⋅dydtanh−1(y)=1−(tanh(x))2pX(x)
(6)获取对数概率密度函数:
l
o
g
P
Y
(
y
)
=
l
o
g
P
X
(
x
)
−
l
o
g
(
1
−
(
t
a
n
h
(
x
)
)
2
)
logP_Y(y) = logP_X(x) - log(1-(tanh(x))^2)
logPY(y)=logPX(x)−log(1−(tanh(x))2)
(7)为了避免
1
−
(
t
a
n
h
(
x
)
)
2
1-(tanh(x))^2
1−(tanh(x))2等于0导致log趋向于无穷小,所以一般会加一个
1
e
−
7
1e-7
1e−7,即:
l
o
g
P
Y
(
y
)
=
l
o
g
P
X
(
x
)
−
l
o
g
(
1
−
(
t
a
n
h
(
x
)
)
2
+
1
e
−
7
)
logP_Y(y) = logP_X(x) - log(1-(tanh(x))^2 + 1e-7)
logPY(y)=logPX(x)−log(1−(tanh(x))2+1e−7)
(8)综上,得到了代码中的结果。