Политический градиент не “обучается”

Я пытаюсь реализовать метод градиента политики из книги “Практическое машинное обучение” автора Жерона, который можно найти здесь. В блокноте используется Tensorflow, а я пытаюсь сделать это с помощью PyTorch.

Мои модели выглядят следующим образом:

model = nn.Sequential(
    nn.Linear(4, 128),
    nn.ELU(),
    nn.Linear(128, 2),
)

Критерий и оптимизаторы:

criterion = nn.BCEWithLogitsLoss()
optim = torch.optim.Adam(model.parameters(), lr=0.01)

Обучение:

env = gym.make("CartPole-v0")

n_games_per_update = 10
n_max_steps = 1000
n_iterations = 250
save_iterations = 10
discount_rate = 0.95

for iteration in range(n_iterations): # Запускаем игру 250 раз
    all_rewards = []
    all_gradients = []
    n_steps = []
    optim.zero_grad()
    for game in range(n_games_per_update): # Запускаем игру 10 раз для накопления градиентов
        current_rewards = []
        current_gradients = []
        obs = env.reset()
        for step in range(n_max_steps): # Запускаем одну игру максимум 1000 шагов

            logit = model(torch.tensor(obs, dtype=torch.float))
            output = F.softmax(logit, dim=0)
            c = Categorical(output)
            action = c.sample()

            y = torch.tensor([1.0 - action, action], dtype=torch.float)
            loss = criterion(logit, y)
            loss.backward()

            obs, reward, done, info = env.step(int(action))
            current_rewards.append(reward)
            current_gradients.append([p.grad for p in model.parameters()])
            if done:
                break
        n_steps.append(step)

        all_rewards.append(current_rewards)
        all_gradients.append(current_gradients)

    # Выполняем дисконтирование и нормализацию
    all_rewards = discount_and_normalize_rewards(all_rewards, discount_rate=discount_rate)

    # Для каждой партии из 10 игр умножаем дисконтированные вознаграждения на градиенты сети. Затем берем среднее для каждого слоя
    new_gradients = []
    for var_index, gradient_placeholder in enumerate(gradient_placeholders):
        means = []
        for game_index, rewards in enumerate(all_rewards):
            for step, reward in enumerate(rewards):
                means.append(reward * all_gradients[game_index][step][var_index])
        new_gradients.append(torch.mean(torch.stack(means), 0, True).squeeze(0))

    # Применяем новые градиенты к сети
    for p, g in zip(model.parameters(), new_gradients):
        p.grad = g.clone()
    optim.step()

Когда я запускаю код на 250 взаимодействиях, я печатаю среднюю длину игры, которую получаю:

Iteration: 50, Average Length: 18.2
Iteration: 100, Average Length: 23.4
Iteration: 150, Average Length: 29.9
Iteration: 200, Average Length: 11.2
Iteration: 250, Average Length: 38.6

Сеть на самом деле не улучшается, и более длительное обучение не помогает. У меня два вопроса:
1. Есть ли что-то очевидно неправильное в том, что я делаю?
2. Я заметил, что в реализации на tensorflow используется логарифм вероятности, но я не уверен, как его интегрировать здесь

Я не могу сказать точно, но думаю, что проблема в том, что вы не вычитаете среднее значение вознаграждений.

Идея заключается в том, что действия с выше среднего вознаграждением становятся положительными после нормализации по среднему, в то время как действия с ниже среднего вознаграждением становятся отрицательными после нормализации по среднему.

Ваш шаг обновления это -log(P(action))*reward, который затем минимизируете с помощью вашего оптимизатора.

P(action)<1 следовательно log(P(action))<0, -log(P(action))>0

Если reward>0, -log(P(action))*reward>0. Минимизация этого значения эквивалентна максимизации log(P(action))*reward<0, что достигается, когда P(action)=1.

С другой стороны, если reward<0, -log(P(action))*reward<0. Это имеет противоположный эффект, при котором P(action) стремится к 0.

Важно то, что разные знаки для вознаграждений выше/ниже среднего вызывают увеличение вероятности действий, связанных с хорошими вознаграждениями, в то время как вероятности действий, связанных с плохими вознаграждениями, уменьшаются.