mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-05-19 21:40:19 +00:00
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from stable_baselines3 import A2C, PPO, SAC, TD3
|
|
from stable_baselines3.common.identity_env import FakeImageEnv
|
|
|
|
SAVE_PATH = './cnn_model.zip'
|
|
|
|
|
|
@pytest.mark.parametrize('model_class', [A2C, PPO, SAC, TD3])
|
|
def test_cnn(model_class):
|
|
# Fake grayscale with frameskip
|
|
# Atari after preprocessing: 84x84x1, here we are using lower resolution
|
|
# to check that the network handle it automatically
|
|
env = FakeImageEnv(screen_height=40, screen_width=40, n_channels=1,
|
|
discrete=model_class not in {SAC, TD3})
|
|
if model_class in {A2C, PPO}:
|
|
kwargs = dict(n_steps=100)
|
|
else:
|
|
# Avoid memory error when using replay buffer
|
|
# Reduce the size of the features
|
|
kwargs = dict(buffer_size=250,
|
|
policy_kwargs=dict(features_extractor_kwargs=dict(features_dim=32)))
|
|
model = model_class('CnnPolicy', env, **kwargs).learn(250)
|
|
|
|
obs = env.reset()
|
|
|
|
action, _ = model.predict(obs, deterministic=True)
|
|
|
|
model.save(SAVE_PATH)
|
|
del model
|
|
|
|
model = model_class.load(SAVE_PATH)
|
|
|
|
# Check that the prediction is the same
|
|
assert np.allclose(action, model.predict(obs, deterministic=True)[0])
|
|
|
|
os.remove(SAVE_PATH)
|