Policy gradient中为什么必须给概率取log？

注意：这篇文章没有解释原因，只是探索了一个相似方法

一般认为是可以用来简化计算的，log可以将原本的累乘转换为累加，例如 a*b*c 总体加上一个log就可以转换为 loga+logb+logc，当然，能总体加log是因为log保留的原来函数的单调方向，转成加法还有个好处就是一定程度上可以防止梯度消失或爆炸，因为连乘很容易为0或无穷

使用最朴素的policy gradient写了一个玩平衡杆的策略模型，发现如果使用SGD会出现一些问题

# 以下代码是有问题的
...
action_probabilities = model(current_env)
...
choose_action_prob = action_probabilities[choose_action]
...
loss = -reward * choose_action_prob
loss.backward()
...

问题出现在直接使用 choose_action_prob 用作梯度上升。

其实直观上感觉应该就是这样的：一局游戏的得分的期望等于这局的total reward乘以这局游戏中各个过程的出现概率。目标就是最大化这个期望，如果使用现有的梯度下降框架，则直接在前面乘以一个负号即可。感觉应该就是这样的，实验结果表明这样模型很有可能不会收敛，模型不会选择最优策略。

我猜测问题可能出现在负号上，理论上这个loss可以无限小，也就是它可能是无解的，但可以通过修改代码使得模型可以趋近于一个值（注：问题原因并不一定是这个，我猜的，但这个解决方法确实可以解决这个问题）

解决方法是让 choose_action_prob = -1/choose_action_prob，因为我们的目的是使得上面的loss更小，也就是使得这里的choose_action_prob更大，前面加上一个负号就可以让它和loss前面的负号抵消，然后让它做除数，即不影响它原本的单调性。

我实验了一下，这种方法确实可行，但更常用的方法是取log

# 方法一
choose_action_prob = 1/choose_action_prob
loss = reward * choose_action_prob
# 方法二，注意后面加一个很小的数防止为0
choose_action_prob = torch.log(choose_action_prob + 1e-8)
loss = -reward * choose_action_prob

这样一来，loss应该是越趋近于0越好。不过实验结果表明，方法二的效果远远好于方法一，甚至于 -1/x+1 也比上面的1/x要好得多，why？

这是 log_e(x) 和 -1/x+1 的曲线

以下为完整代码：

import itertools
import threading

import gymnasium
import numpy as np
import torch
from loguru import logger
from torch.utils.tensorboard import SummaryWriter

board = SummaryWriter('./logs/division_factor_random')

torch.random.manual_seed(1)
np.random.seed(0)


def visualize_train(policy: torch.nn.Module):
    with torch.no_grad():
        test_env = gymnasium.make("CartPole-v1", render_mode='human')
        for i in itertools.count():
            total_reward = 0
            obs, _ = test_env.reset(seed=0)
            while True:
                action = torch.argmax(policy(torch.FloatTensor(obs)))
                obs, reward, terminated, _, info = test_env.step(action.item())
                total_reward += reward
                if terminated:
                    board.add_scalar("test/reward", total_reward, i)
                    logger.info(f"total reward: {total_reward}")
                    break
        test_env.close()


class Actor(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.dnn = torch.nn.Sequential(
            torch.nn.Linear(in_features=4, out_features=20),
            torch.nn.ReLU(),
            torch.nn.Linear(in_features=20, out_features=2),
            torch.nn.Softmax(dim=0)
        )

    def forward(self, x):
        return self.dnn(x)


actor = Actor()

threading.Thread(target=visualize_train, args=(actor,)).start()

optimizer = torch.optim.Adam(params=actor.parameters())

env = gymnasium.make("CartPole-v1", render_mode=None)

for game_count in itertools.count():
    obs, _ = env.reset(seed=0)

    acts = []
    rewards = []
    for s in itertools.count():
        action_probs = actor(torch.tensor(obs, dtype=torch.float))
        action = torch.distributions.Categorical(action_probs).sample().item()
        # if np.random.random() < 0.1: action = np.random.randint(0, 2)

        obs, reward, terminated, _, info = env.step(action)

        acts.append(1 / action_probs[action])
        rewards.append(reward)

        if terminated:
            # loss with discount factor=0.99
            loss = sum([a * weighted_r for a, weighted_r in zip(acts, [sum([r * 0.99 ** j for j, r in enumerate(rewards[i:])]) for i in range(len(rewards))])])
            # loss = (sum(rewards) * sum(acts))
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            break

env.close()

Leave a Comment Cancel reply