Обучение трансформера с помощью PPO

Контекст

Я пытаюсь применить обучение с подкреплением к трансформеру.

У меня есть следующие токены: ["<sos>", "<eos>", "roses", "are", "red"].

Моя цель – использовать основанную на трансформере политику сети для генерации токенов ["<sos>", "red", "<eos>"] на основе ввода ["roses", "are"].

Сеть политики принимает исходные и целевые входные данные для генерации логитов для следующего токена.

Например, исходный вход ["roses", "are"], а целевой вход ["<sos>"].

С этими входными данными сеть политики будет генерировать логиты, которые, надеюсь, будут соответствовать токену "red".

Я повторяю этот процесс в авторегрессивном режиме, пока не будет сгенерирован токен "<eos>".

Для сети ценностей я использую LSTM.

Сеть ценностей принимает входные данные исходных и целевых токенов так же, как и трансформер.

Она объединяет исходные и целевые токены и генерирует одно значение (скрытое состояние), представляющее ценность.

Например, она объединяет исходные ["roses", "are"] и целевые ["<sos>", "red", "<eos>"] в ["roses", "are", "<sos>", "red", "<eos>"] и генерирует число.

Я обучаю сети политики и ценности, используя PPO и GAE.

Каждый эпизод завершается, когда сеть политики заканчивает генерировать токены.

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

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

У меня есть следующие правила вознаграждений/наказаний:

1. Добавить -1 к счету за каждую несоответствующую длину целевого токена. Например, желаемая длина целевого токена – 3 (["<sos>", "red", "<eos>"]). Если политика сгенерировала 5 токенов, будет добавлено -2 к счету в наказание за неправильную генерацию.
2. Добавить +1 к счету за каждый правильный токен с правильным индексом.
3. Добавить +10 к счету, если сгенерированные токены exactly совпадают с тем, что я ожидал.

Вопрос

Результат стабилизируется с неправильными значениями во время обучения.

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

Тем не менее, она неверно изменила свое поведение в последующих эпизодах.

Я хочу знать, что я делаю неправильно.

Это просто ошибка программирования, неправильные правила вознаграждений или недопонимание некоторых теорий?

Вот код, который я написал, используя PyTorch:

import torch
import torch.nn as nn
from tensordict import TensorDict
from tensordict.nn import TensorDictModule, InteractionType
from torch import optim
from torch.distributions import Categorical
from torch.nn import Embedding
from torchrl.modules import ProbabilisticActor, ValueOperator
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value.functional import generalized_advantage_estimate

class OurEmbedding:

    def __init__(self, device):
        self.token_to_index = {
            "<sos>": 0,
            "<eos>": 1,
            "roses": 2,
            "are": 3,
            "red": 4
        }
        self.vocab_size = len(self.token_to_index)
        self.token_vec_dim = 512
        self.device = device
        self.torch_embedding = Embedding(self.vocab_size, self.token_vec_dim, device=self.device)

class PolicyNetwork(nn.Module):

    def __init__(self, our_embedding, device):
        super().__init__()
        self.our_embedding = our_embedding
        self.device = device
        self.transformer = nn.Transformer(
            d_model=self.our_embedding.token_vec_dim,
            nhead = 8 if self.our_embedding.token_vec_dim % 8 == 0 else 2,
            batch_first = True,
            device=self.device
        )
        self.fc = nn.Sequential(
            nn.Linear(self.our_embedding.token_vec_dim, self.our_embedding.token_vec_dim * 2, device=self.device),
            nn.ReLU(),
            nn.Linear(self.our_embedding.token_vec_dim * 2, self.our_embedding.vocab_size, device=self.device)
        )

    def forward(self, src_token_vecs, target_token_vecs):
        """
        Предсказать следующий токен.
        Обратите внимание, что мы предполагаем, что входные данные не пакетированные.
        :param src_token_vecs: То же самое, что и `nn.Transformer.forward(src, tgt)`.
                               Размер: (число_токенов, размер_вектора_токена).
        :param target_token_vecs: То же самое, что и `nn.Transformer.forward(src, tgt)`.
                                  Размер: (число_токенов, размер_вектора_токена).
        :return: Сырой счет (логит) для следующего токена.
                 Размер: (размер_словаря).
        """
        # Размер: (число_токенов, размер_вектора_токена).
        transformer_output = self.transformer(src_token_vecs, target_token_vecs)
        # Размер: (размер_вектора_токена).
        last_token_vec = transformer_output[-1]
        return self.fc(last_token_vec)

class ValueNetwork(nn.Module):

    def __init__(self, our_embedding, device):
        super().__init__()
        self.our_embedding = our_embedding
        self.device = device
        hidden_size = self.our_embedding.token_vec_dim * 2
        self.lstm = nn.LSTM(
            input_size=self.our_embedding.token_vec_dim,
            hidden_size=hidden_size,
            device=self.device
        )
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, hidden_size * 2, device=self.device),
            nn.ReLU(),
            nn.Linear(hidden_size * 2, 1, device=self.device)
        )

    def forward(self, src_token_vecs, generated_token_vecs):
        """
        Оценить производительность политики.
        Обратите внимание, что мы предполагаем, что входные данные не пакетированные.
        :param src_token_vecs: То же самое, что и `nn.Transformer.forward(src, tgt)`.
                               Размер: (число_токенов, размер_вектора_токена).
        :param generated_token_vecs: Размер: (число_токенов, размер_вектора_токена).
        :return: Ценность. Размер: (1).
        """
        # Размер: (число_объединенных_токенов, размер_вектора_токена).
        input_tensor = torch.cat((src_token_vecs, generated_token_vecs), dim=0).to(self.device)
        # Размер `hidden_state`: (1, hidden_size).
        # Мы предполагаем, что LSTM однонаправленный, и количество слоев равно 1.
        _, (hidden_state, _) = self.lstm(input_tensor)
        fc_output = self.fc(hidden_state.squeeze(dim=0))
        return self.fc(hidden_state.squeeze(dim=0))

def main():
    seed = 3
    device = torch.device("cpu")
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        device = torch.device("cuda")
        torch.cuda.manual_seed_all(seed)

    our_embedding = OurEmbedding(device)
    policy_network = PolicyNetwork(our_embedding, device)
    policy_module = TensorDictModule(
        module=policy_network,
        in_keys=["src_token_vecs", "target_token_vecs"],
        out_keys=["logits"]
    )
    actor = ProbabilisticActor(
        module=policy_module,
        in_keys=["logits"],
        # Размер: скалярное значение, представляющее индекс следующего токена.
        out_keys=["action"],
        distribution_class=Categorical,
        # Принимать решение в стохастическом режиме.
        default_interaction_type=InteractionType.RANDOM,
        return_log_prob=True
    )
    value_network = ValueNetwork(our_embedding, device)
    value_operator = ValueOperator(
        module=value_network,
        in_keys=["src_token_vecs", "target_token_vecs"],
        out_keys=["value"]
    )
    loss_module = ClipPPOLoss(
        actor_network=actor,
        critic_network=value_operator,
    )
    loss_module.set_keys(
        advantage="advantage",
        value_target="value_target",
        value="value",
        action="action",
        sample_log_prob="sample_log_prob"
    )

    # Обучение

    episode = 0
    max_episode = 1000
    optimizer = optim.Adam(list(policy_network.parameters()) + list(value_network.parameters()))
    while episode < max_episode:
        # Мы начинаем с 1, потому что начальный токен "<sos>" уже использован.
        token_gen_iteration = 1
        done = False
        # Он будет генерировать максимум 10 токенов.
        max_token_gen_iteration = 10
        # Размер: (число_токенов, размер_вектора_токена).
        src_token_vecs = our_embedding.torch_embedding(torch.LongTensor([2, 3]).to(device))
        # Предполагая, что индекс=0 соответствует токену "<sos>".
        target_token_indexes = [0]
        # Это изменяемое.
        # Размер: (число_токенов, размер_вектора_токена).
        target_token_vecs = our_embedding.torch_embedding(torch.LongTensor(target_token_indexes).to(device))
        # Это изменяемое.
        current_tensor_dict = TensorDict({
            "src_token_vecs": src_token_vecs,
            "target_token_vecs": target_token_vecs
        }, batch_size=[])
        while not done and token_gen_iteration < max_token_gen_iteration:
            current_tensor_dict = TensorDict({
                "src_token_vecs": src_token_vecs,
                "target_token_vecs": target_token_vecs
            }, batch_size=[])
            actor(current_tensor_dict)
            # Размер: скаляр.
            next_token_index = current_tensor_dict["action"]
            target_token_indexes.append(next_token_index.item())
            # Исходная форма (размер_вектора_токена).
            # Поскольку `target_token_vecs` имеет размер (число_токенов, размер_вектора_токена),
            # нам нужно добавить еще одно измерение, чтобы размер стал (1, размер_вектора_токена) для операции `torch.cat`.
            next_token_vec = our_embedding.torch_embedding(next_token_index).unsqueeze(dim=0)
            # Поскольку мы хотим добавить `next_token_vec` как дополнительный целевой токен, размер `torch.cat` - 0.
            # Размер: (число_токенов++, то же самое).
            target_token_vecs = torch.cat((target_token_vecs, next_token_vec), dim=0)
            # Предполагая, что индекс=1 соответствует токену "<eos>".
            if next_token_index.item() == 1:
                done = True
            token_gen_iteration += 1
        current_tensor_dict["sample_log_prob"] = current_tensor_dict["sample_log_prob"].detach()
        value_operator(current_tensor_dict)
        next_tensor_dict = TensorDict({
            "src_token_vecs": src_token_vecs,
            "target_token_vecs": target_token_vecs
        }, batch_size=[])
        value_operator(next_tensor_dict)

        # Правила вознаграждений/наказаний
        score = 0
        # Наказание за неправильное количество токенов.
        score -= abs(len(target_token_indexes) - 3)
        # Вознаграждение за правильные индексы.
        if len(target_token_indexes) > 1 and target_token_indexes[1] == 4:
            score += 1
        if len(target_token_indexes) > 2 and target_token_indexes[2] == 1:
            score += 1
        if target_token_indexes == [0, 4, 1]:
            score += 10
        reward = torch.FloatTensor([[score]]).to(device)

        # Обратите внимание, что нам нужно использовать пакетированный ввод, и выход будет в пакетированной форме.
        advantage, value_target = generalized_advantage_estimate(
            gamma=0.98,
            lmbda=0.95,
            state_value=current_tensor_dict["value"].unsqueeze(0),
            next_state_value=next_tensor_dict["value"].unsqueeze(0),
            reward=reward,
            done=torch.BoolTensor([[1]]).to(device),
            terminated=torch.BoolTensor([[1]]).to(device)
        )
        current_tensor_dict["advantage"] = advantage.squeeze(0)
        current_tensor_dict["value_target"] = value_target.squeeze(0)
        loss_tensor_dict = loss_module(current_tensor_dict)
        loss_critic = loss_tensor_dict["loss_critic"]
        loss_entropy = loss_tensor_dict["loss_entropy"]
        loss_objective = loss_tensor_dict["loss_objective"]
        loss = loss_critic + 0.01 * loss_entropy + loss_objective
        print(f"episode: {episode}, loss_critic: {loss_critic}, loss_objective: {loss_objective}, loss_entropy: {loss_entropy}, score: {score}, target_token_indexes: {target_token_indexes}")
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        episode += 1

main()

И вот фрагмент вывода консоли:

episode: 0, loss_critic: 6.493023872375488, loss_objective: 5.944887638092041, loss_entropy: -0.015881547704339027, score: -7, target_token_indexes: [0, 0, 4, 3, 3, 4, 3, 0, 3, 3]
episode: 1, loss_critic: 5.89158821105957, loss_objective: 6.10428524017334, loss_entropy: -0.014061596244573593, score: -7, target_token_indexes: [0, 0, 0, 0, 4, 4, 4, 2, 0, 0]
episode: 2, loss_critic: 1.898949384689331, loss_objective: -2.878739356994629, loss_entropy: -0.012305832467973232, score: -1, target_token_indexes: [0, 1]
episode: 3, loss_critic: 18.447572708129883, loss_objective: -15.853285789489746, loss_entropy: -0.012489358894526958, score: 12, target_token_indexes: [0, 4, 1]
episode: 4, loss_critic: 8.048697471618652, loss_objective: -8.492928504943848, loss_entropy: -0.00397573271766305, score: -1, target_token_indexes: [0, 1]
episode: 5, loss_critic: 9.843278884887695, loss_objective: -10.317561149597168, loss_entropy: -0.0009762179106473923, score: -1, target_token_indexes: [0, 1]
episode: 6, loss_critic: 11.368447303771973, loss_objective: -11.865358352661133, loss_entropy: -0.0002861745888367295, score: -1, target_token_indexes: [0, 1]
episode: 7, loss_critic: 12.705635070800781, loss_objective: -13.205760955810547, loss_entropy: -8.256323781097308e-05, score: -1, target_token_indexes: [0, 1]
episode: 8, loss_critic: 13.895231246948242, loss_objective: -14.395499229431152, loss_entropy: -3.16958139592316e-05, score: -1, target_token_indexes: [0, 1]
episode: 9, loss_critic: 14.960994720458984, loss_objective: -15.460846900939941, loss_entropy: -1.4596819710277487e-05, score: -1, target_token_indexes: [0, 1]

Как видно, он однажды сгенерировал идеальный результат ([0, 4, 1]).

Тем не менее, модель изменила свое поведение и выдала [0, 1] и закрепилась на этом значении.

Я нашел свои ошибки и исправил их.

Проблема 1: Недостаток исследования

Время оптимизации было плохим.

Я думал, что отрицательное вознаграждение уменьшит шанс выбора этого действия и наоборот.

Однако я осознал, что абсолютная величина вознаграждения не имеет значения.

Вместо этого относительная разница между вознаграждениями приводит к желаемым изменениям.

Ранее я оптимизировал агента, как только он получал вознаграждение.

Это только побуждало агента делать то же самое.

Поэтому я отложил оптимизацию каждые 100 эпизодов, чтобы агент мог испытать различные результаты и соответствующие вознаграждения.

Проблема 2: Отброшенные промежуточные действия

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

Однако я должен был итерировать историю действий и вычислять вознаграждения для каждого действия.

Проблема 3: Неэффективная сеть ценностей

Ранее я использовал LSTM для вычисления ценности.

Я заметил, что использование трансформера для сети ценностей значительно улучшило производительность.

Честно говоря, я не знаю, почему основанная на LSTM сеть ценностей не работала хорошо.

Решение

Вот обновленный код:

import torch
import torch.nn as nn
from tensordict import TensorDict
from tensordict.nn import TensorDictModule, InteractionType
from torch import optim
from torch.distributions import Categorical
from torch.nn import Embedding
from torchrl.modules import ProbabilisticActor, ValueOperator
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value.functional import generalized_advantage_estimate

class OurEmbedding:

    def __init__(self, device):
        self.token_to_index = {
            "<sos>": 0,
            "<eos>": 1,
            "roses": 2,
            "are": 3,
            "red": 4
        }
        self.vocab_size = len(self.token_to_index)
        self.token_vec_dim = 8
        self.device = device
        self.torch_embedding = Embedding(self.vocab_size, self.token_vec_dim, device=self.device)

class PolicyNetwork(nn.Module):

    def __init__(self, our_embedding, device):
        super().__init__()
        self.our_embedding = our_embedding
        self.device = device
        self.transformer = nn.Transformer(
            d_model=self.our_embedding.token_vec_dim,
            nhead = 8 if self.our_embedding.token_vec_dim % 8 == 0 else 2,
            batch_first = True,
            device=self.device
        )
        self.fc = nn.Sequential(
            nn.Linear(self.our_embedding.token_vec_dim, self.our_embedding.token_vec_dim * 2, device=self.device),
            nn.ReLU(),
            nn.Linear(self.our_embedding.token_vec_dim * 2, self.our_embedding.vocab_size, device=self.device)
        )

    def forward(self, src_token_vecs, target_token_vecs):
        """
        Предсказать следующий токен.
        Обратите внимание, что мы предполагаем, что входные данные не пакетированные.
        :param src_token_vecs: То же самое, что и `nn.Transformer.forward(src, tgt)`.
                               Размер: (число_токенов, размер_вектора_токена).
        :param target_token_vecs: То же самое, что и `nn.Transformer.forward(src, tgt)`.
                                  Размер: (число_токенов, размер_вектора_токена).
        :return: Сырой счет (логит) для следующего токена.
                 Размер: (размер_словаря).
        """
        # Размер: (число_токенов, размер_вектора_токена).
        transformer_output = self.transformer(src_token_vecs, target_token_vecs)
        # Размер: (размер_вектора_токена).
        last_token_vec = transformer_output[-1]
        return self.fc(last_token_vec)

class ValueNetwork(nn.Module):

    def __init__(self, our_embedding, device):
        super().__init__()
        self.our_embedding = our_embedding
        self.device = device
        self.transformer = nn.Transformer(
            d_model=self.our_embedding.token_vec_dim,
            nhead=8 if self.our_embedding.token_vec_dim % 8 == 0 else 2,
            batch_first=True,
            device=self.device
        )
        self.fc = nn.Sequential(
            nn.Linear(self.our_embedding.token_vec_dim, self.our_embedding.token_vec_dim * 2, device=self.device),
            nn.ReLU(),
            nn.Linear(self.our_embedding.token_vec_dim * 2, 1, device=self.device)
        )

    def forward(self, src_token_vecs, generated_token_vecs):
        """
        Оценить производительность политики.
        Обратите внимание, что мы предполагаем, что входные данные не пакетированные.
        :param src_token_vecs: То же самое, что и `nn.Transformer.forward(src, tgt)`.
                               Размер: (число_токенов, размер_вектора_токена).
        :param generated_token_vecs: Размер: (число_токенов, размер_вектора_токена).
        :return: Ценность. Размер: (1).
        """
        # Размер: (число_токенов, размер_вектора_токена).
        transformer_output = self.transformer(src_token_vecs, generated_token_vecs)
        last_element = transformer_output[-1]
        return self.fc(last_element)

def main():
    torch.autograd.set_detect_anomaly(True)
    seed = 3
    device = torch.device("cpu")
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        device = torch.device("cuda")
        torch.cuda.manual_seed_all(seed)

    our_embedding = OurEmbedding(device)
    policy_network = PolicyNetwork(our_embedding, device)
    policy_module = TensorDictModule(
        module=policy_network,
        in_keys=["src_token_vecs", "target_token_vecs"],
        out_keys=["logits"]
    )
    actor = ProbabilisticActor(
        module=policy_module,
        in_keys=["logits"],
        # Размер: скалярное значение, представляющее индекс следующего токена.
        out_keys=["action"],
        distribution_class=Categorical,
        # Принимать решение в стохастическом режиме.
        default_interaction_type=InteractionType.RANDOM,
        return_log_prob=True
    )
    value_network = ValueNetwork(our_embedding, device)
    value_operator = ValueOperator(
        module=value_network,
        in_keys=["src_token_vecs", "target_token_vecs"],
        out_keys=["value"]
    )
    loss_module = ClipPPOLoss(
        actor_network=actor,
        critic_network=value_operator
    )
    loss_module.set_keys(
        advantage="advantage",
        value_target="value_target",
        value="value",
        action="action",
        sample_log_prob="sample_log_prob"
    )

    # Обучение

    epoch = 0
    max_epoch = 10000
    optimizer = optim.Adam(loss_module.parameters())
    total_loss = 0
    while epoch < max_epoch:
        # Мы начинаем с 1, потому что начальный токен "<sos>" уже использован.
        token_gen_iteration = 1
        # Он будет генерировать максимум 10 токенов.
        max_token_gen_iteration = 10
        # Размер: (число_токенов, размер_вектора_токена).
        src_token_vecs = our_embedding.torch_embedding(torch.LongTensor([2, 3]).to(device))
        # Предполагая, что индекс=0 соответствует токену "<sos>".
        target_token_indexes = [0]
        # Это изменяемое.
        tensor_dict = TensorDict({}, batch_size=[])
        current_tensor_dict = tensor_dict
        # Принимать действия, пока не получим токен "<eos>" или не достигнем максимальной длины токентов.
        while True:
            # Размер: (число_токенов, размер_вектора_токена).
            target_token_vecs = our_embedding.torch_embedding(torch.LongTensor(target_token_indexes).to(device))
            current_tensor_dict["src_token_vecs"] = src_token_vecs
            current_tensor_dict["target_token_vecs"] = target_token_vecs
            value_operator(current_tensor_dict)
            # Предполагая, что индекс=1 соответствует токену "<eos>".
            if target_token_indexes[-1] == 1 or token_gen_iteration >= max_token_gen_iteration:
                break
            actor(current_tensor_dict)
            current_tensor_dict["sample_log_prob"] = current_tensor_dict["sample_log_prob"].detach()
            # Размер: скаляр.
            next_token_index = current_tensor_dict["action"]
            target_token_indexes.append(next_token_index.item())
            current_tensor_dict["next"] = TensorDict({}, batch_size=[])
            current_tensor_dict = current_tensor_dict["next"]
            token_gen_iteration += 1
        # Итерируемся по tensor_dict и вычисляем преимущества.
        current_tensor_dict = tensor_dict
        while "next" in current_tensor_dict:
            done = torch.BoolTensor([[0]]).to(device)
            terminated = torch.BoolTensor([[0]]).to(device)
            if current_tensor_dict["action"].item() == 1:
                done = torch.BoolTensor([[1]]).to(device)
                terminated = torch.BoolTensor([[1]]).to(device)

            # Правила вознаграждений/наказаний
            score = 0
            if len(target_token_indexes) > 1 and target_token_indexes[1] == 4:
                score += 3
            if len(target_token_indexes) == 3:
                score += 5
            if target_token_indexes == [0, 4, 1]:
                score += 10

            next_tensor_dict = current_tensor_dict["next"]
            # Обратите внимание, что нам нужно использовать пакетированный ввод, и выход будет в пакетированной форме.
            advantage, value_target = generalized_advantage_estimate(
                gamma=0.98,
                lmbda=0.95,
                state_value=current_tensor_dict["value"].unsqueeze(0),
                next_state_value=next_tensor_dict["value"].unsqueeze(0),
                reward=torch.FloatTensor([[score]]).to(device),
                done=done,
                terminated=terminated
            )
            current_tensor_dict["advantage"] = advantage.squeeze(0)
            current_tensor_dict["value_target"] = value_target.squeeze(0)
            current_tensor_dict = current_tensor_dict["next"]
        # Итерируемся по tensor_dict и вычисляем каждую потерю.
        current_tensor_dict = tensor_dict
        while "next" in current_tensor_dict:
            loss_tensor_dict = loss_module(current_tensor_dict)
            loss_critic = loss_tensor_dict["loss_critic"]
            loss_entropy = loss_tensor_dict["loss_entropy"]
            loss_objective = loss_tensor_dict["loss_objective"]
            loss = loss_critic + loss_entropy + loss_objective
            total_loss += loss
            current_tensor_dict = current_tensor_dict["next"]
        if epoch % 100 == 0:
            total_loss.backward()
            optimizer.step()
            optimizer.zero_grad()
            total_loss = 0
        print(f"epoch: {epoch}, target_token_indexes: {target_token_indexes}")
        epoch += 1

main()

Кстати, обучение занимает много времени.

Я думаю, что ждал около 12 часов.

По крайней мере, в конце он выучил правильные действия 🙂