Немного запутался с обучение LunarLander на PyTorch и Google Colab

Начинаю изучать RL. Где-то ошибка. но все выглядит верно. Буду благодарен любому совету спасибо. Заранее спасибо

https://colab.research.google.com/drive/1SqcNoRgu0sv2KFLSSYs4wX7mp49xJpou?usp=sharing

# установка необходимых билиотек
!apt-get install -y xvfb python-opengl > /dev/null 2>&1
!pip install colabgymrender
!pip install imageio==2.4.1
!pip install gym pyvirtualdisplay > /dev/null 2>&1
!pip install gym[box2d]

# импортирование библиотек
import gym
from gym.wrappers.time_limit import TimeLimit
from colabgymrender.recorder import Recorder
import torch
from torch import nn, optim
from torch.optim.adam import Adam
import numpy as np

# создание переменных
DIRECTORY: str = './video'
DEVICE: str = "cuda" if torch.cuda.is_available() else "cpu"
EPOCHS: int = 1000

# создание среды
env: Recorder = Recorder(
    gym.make("LunarLander-v2"), 
    DIRECTORY
)

# объявление класса
class Lander(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.__model = nn.Sequential(
            nn.Linear(in_features=8, out_features=16).type(torch.float32),
            nn.Linear(in_features=16, out_features=8).type(torch.float32),
            nn.Linear(in_features=8, out_features=4).type(torch.float32),
        )
        
    def forward(self, X: torch.Tensor) -> int:
        '''
        Метод forward propagation
        :params: torch.Tensor: X: входной тензор
        '''
        # вызов метода forward propagation
        X: torch.Tensor = self.__model(X)
         
        # выбор текущего действия
        X: int = torch.argmin(X).item()

        # возвращение индекса действия
        return X

    @staticmethod
    def start_landing_procedure(X: torch.Tensor) -> torch.Tensor:
        '''
        объединение лосс функции и forward в одной функции,
        стек помечается на удаление СМ после полной попытки посадить корабль
        :params: torch.Tensor: X: входной вектор
        :return: результат,  данном случае, reward
        '''
        # так как мы получаем результат посадки, то следующая эпоха - это новая попытка посадить корабль
        # сброс среды
        X: np.ndarray = env.reset()

        # перевод массива np.ndarray в массив torch.Tensor
        X: torch.Tensor = torch.from_numpy(X).to(DEVICE)

        # проверка закончена ли текущая попытка
        TERMINAL: bool = False

        # создание начальной ошибки
        TOTAL_REWARD: float = 0.0

        # проигрываем эпизод
        while not TERMINAL:
            # получаем следующее действие с помощью текущих параметров нейронной сети
            ACTION: int = lander.forward(X)

            # совершение следующего шага с помощь виртуальной среды opengym AI
            X, REWARD, TERMINAL, _ = env.step(ACTION)

            # перевод наблюдения в массив тензор для передачи в нейронную сеть
            X: torch.Tensor = torch.from_numpy(X).to(DEVICE)

            # подсчёт ошибок на каждом шаге, чтобы получить общую ошибку
            # попытки посадить модуль на Луну
            TOTAL_REWARD += REWARD
        
        # возвращаем награду, фактически - это ошибка
        TOTAL_REWARD: torch.Tensor = torch.tensor(float(TOTAL_REWARD), requires_grad=True).type(torch.float32)

        # нас интересует gradient ascent, поэтому результат возвращаем со знаком минус
        return -TOTAL_REWARD


lander: Lander = Lander().to(DEVICE)
optimizer: Adam = optim.Adam(lander.parameters(), lr=0.0001)

# перевод модели в режим обучения
# тестирование модели не нужно
lander.train()

# так как мы получаем результат посадки, то следующая эпоха - это новая попытка посадить корабль
# сброс среды
X: np.ndarray = env.reset()

# перевод массива np.ndarray в массив torch.Tensor
X: torch.tensor = torch.from_numpy(X).to(DEVICE)

# запуск цикла с обучение, где кол-во эпох - это кол-во посадить корабль
for EPOCH in range(EPOCHS):

    # получаем результат посадки
    Landing_error: torch.Tensor = lander.start_landing_procedure(X)

    # обнуление градиента с предыдущего шага
    optimizer.zero_grad()
    
    # расчёт градиента каждого элемента (матрица Гессе, если не ошибаюсь :-))
    Landing_error.backward()

    # смещение значений с помощью оптимизацтора Adam с использованием градиента
    optimizer.step()

    LANDING_SCORE: float = round(-Landing_error.item(), 2)

    print(f"epoch: {EPOCH}, error: {LANDING_SCORE}")

нет ощущения, что нейронка обучается((