Агент SB3 PPO не знает о среде ⇐ Python

[Anonymous]

Я пытаюсь создать PPO со стабильными базовыми показателями 3 и пользовательским cnn, который будет играть в Марио, но мой агент сразу же попадает в яму даже после пары часов обучения. Он ведет себя так, как будто он ничего видит.
Мои наблюдения за окружением относятся к (19, 19, 28), где (19, 19) — размер сетки вокруг Марио (9 вверху, 9 справа и т. д.), а 28 — 7 каналов x 4 кадра (складывается с помощью VecFrameStack). Семь каналов представляют собой единое представление каждого типа ячеек, таких как сплошные блоки, топтаемые враги и т. д.
Я пробовал настроить гиперпараметры и функции оценки, но они не сработали, и я подозреваю, что проблема заключается в извлечении признаков.

Буду признателен за любую помощь.
Сценарий обучения:
def make_env(rank):
    def _init():
        env = MarioGymEnv(port=5555+rank)
        env = ThrottleEnv(env, delay=0)
        env = SkipEnv(env, skip=2) # custom environment to skip every other frame
        return env
    return _init

def main():
    num_cpu = 12
    env = SubprocVecEnv([make_env(i) for i in range(num_cpu)])
    env = VecFrameStack(env, n_stack=4)
    env = VecMonitor(env)
    policy_kwargs = dict(
        features_extractor_class=Cnn,
    )
   
    model = PPO(
        'CnnPolicy',
        env,
        policy_kwargs=policy_kwargs,
        verbose=1,
        tensorboard_log='./board',
        learning_rate=1e-3,
        n_steps=256,
        batch_size=256,
    )
    TOTAL_TIMESTEPS = 5_000_000
    TB_LOG_NAME = 'PPO-CustomCNN-ScheduledLR'

    checkpoint_callback = CheckpointCallback(
        save_freq= max(10_000 // num_cpu, 1),
        save_path='./models/',
        name_prefix='marioAI'
    )
   
    try:
        model.learn(
            total_timesteps=TOTAL_TIMESTEPS,
            callback=checkpoint_callback,
            tb_log_name=TB_LOG_NAME
        )
        model.save('marioAI_final')

    except Exception as e:
        print(e)
        model.save('marioAI_error')'

Вот мой экстрактор функций.
class Cnn(BaseFeaturesExtractor):
    def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 256):
        super().__init__(observation_space, features_dim)
        n_input_channels = observation_space.shape[2]
       
        self.cnn = nn.Sequential(
            nn.Conv2d(n_input_channels, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
           
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # Stride 2 downsamples
            nn.ReLU(),
           
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # Stride 2 downsamples
            nn.ReLU(),
        )
       
        with torch.no_grad():
            dummy_input = torch.zeros(
                (1, n_input_channels, observation_space.shape[0], observation_space.shape[1])
            )
           
            output = self.cnn(dummy_input)
            n_flattened_features = output.flatten(1).shape[1]

        self.linear_head = nn.Sequential(
            nn.Linear(n_flattened_features, features_dim),
            nn.ReLU()
        )

    def forward(self, observations: torch.Tensor) -> torch.Tensor:
        observations = observations.permute(0, 3, 1, 2)
        cnn_output = self.cnn(observations)
        flattened_features = torch.flatten(cnn_output, start_dim=1)
        features = self.linear_head(flattened_features)
       
        return features


Подробнее здесь: https://stackoverflow.com/questions/798 ... nvironment