Почему я не могу установить seed для своей программы DQN, используя sbx?

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

Вот попытка создать минимальный воспроизводимый пример –

import gymnasium as gym
import torch
import numpy as np
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.monitor import Monitor
import os
import argparse
import random
from sbx import DQN
def run_trials():

    parser = argparse.ArgumentParser()
    parser.add_argument('--seed', type=int, default=1, help='Случайное начальное значение для воспроизводимости. Если не указано, будет сгенерировано случайное начальное значение.')
    args = parser.parse_args()
    # Если начальное значение не указано, сгенерировать случайное
    if args.seed is None:
        args.seed = random.randint(1, 1000000)
    print(f"Используется начальное значение: {args.seed}")
    # Заменить раздел загрузки конфигурации на эту захардкоженную конфигурацию
    config = {
        "n_timesteps": int(5e4),
        "policy": 'MlpPolicy',
        "learning_rate": 2.3e-3,
        "batch_size": 64,
        "buffer_size": 100000,
        "learning_starts": 1000,
        "gamma": 0.99,
        "target_update_interval": 10,
        "train_freq": 256,
        "gradient_steps": 128,
        "exploration_fraction": 0.16,
        "exploration_final_eps": 0.04,
        "policy_kwargs": dict(net_arch=[256, 256])
    }
    def set_seed(seed):
        """Установить начальное значение для воспроизводимости."""
        os.environ['PYTHONASHSEED'] = str(seed)
        np.random.seed(seed)
        random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    # Установить начальное значение
    set_seed(args.seed)
    def create_env(seed=None, is_eval=False, initial_states=None):
        # Создать базовую среду.
        env = gym.make("CartPole-v1")
        env = Monitor(env)

        # Установить начальные значения, чтобы обеспечить детерминированное поведение.
        if seed is not None:
            env.observation_space.seed(seed)
            env.action_space.seed(seed)
            env.reset(seed=seed)
        return env
    # Создать среды
    env = DummyVecEnv([lambda: create_env(args.seed)])
    policy = config.pop('policy', 'MlpPolicy')
    n_timesteps = config.pop('n_timesteps', int(1e5))  # Использовать get() с значением по умолчанию
    # Создать модель
    model = DQN(
        policy=policy,
        env=env,
        seed=args.seed,
        **config,
    )
    callback = ([])
    model.learn(total_timesteps=int(n_timesteps // 5), callback=callback,
                reset_num_timesteps=False)

    def create_eval_env(seed=None):
        # Создать базовую среду для оценки
        eval_env = gym.make("CartPole-v1")
        # Применить обертку Monitor для отслеживания наград за эпизоды и статистики
        eval_env = Monitor(eval_env)
        # Установить начальные значения, чтобы обеспечить детерминированное поведение
        if seed is not None:
            eval_env.observation_space.seed(seed)
            eval_env.action_space.seed(seed)
            eval_env.reset(seed=seed)
        return eval_env

    def evaluate_model_deterministically(model, seed, num_episodes=100):
        # Создать новую среду для оценки
        eval_env = create_eval_env(seed=seed)
        episode_rewards = []
        for _ in range(num_episodes):
            obs, _ = eval_env.reset()
            done = False
            episode_reward = 0
            while not done:
                action, _ = model.predict(obs, deterministic=True)
                obs, reward, terminated, truncated, _ = eval_env.step(action)
                done = terminated or truncated
                episode_reward += reward
            episode_rewards.append(episode_reward)
        eval_env.close()
        mean_reward = np.mean(episode_rewards)
        std_reward = np.std(episode_rewards)

        return mean_reward, std_reward
    mean_reward, std_reward = evaluate_model_deterministically(model, args.seed)
    return mean_reward, std_reward
mean_rewards = []
std_rewards = []
for run in range(3):
    print("run = ", run)
    mean_reward, std_reward = run_trials()
    mean_rewards.append(mean_reward)
    std_rewards.append(std_reward)

print("mean_rewards = ", mean_rewards)

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

Вот мои результаты –
Попытка 1:

run =  0
Используется начальное значение: 1
run =  1
Используется начальное значение: 1
run =  2
Используется начальное значение: 1
mean_rewards =  [120.52, 120.52, 120.52]

Попытка 2:

run =  0
Используется начальное значение: 1
run =  1
Используется начальное значение: 1
run =  2
Используется начальное значение: 1
mean_rewards =  [116.64, 116.64, 116.64]

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

Я изучил документацию по настройке начального значения для среды здесь и также прочитал следующее: “Совершенно воспроизводимые результаты не гарантируются при различных выпусках PyTorch или на разных платформах. Более того, результаты не обязательно должны быть воспроизводимыми при выполнении на CPU и GPU, даже при использовании идентичных начальных значений.“. Однако я хотел бы убедиться, что на моей стороне нет ошибок. Также я использую sbx вместо stable-baselines3. Возможно, это проблема JAX?