saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-27 20:02:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add test for predict method

Browse source
This commit is contained in:
Antonin Raffin 2020-02-14 14:03:41 +01:00
parent 8b559d71ab
commit e31b139c47
4 changed files with 75 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
tests/test_predict.py Normal file
View file

@ -0,0 +1,55 @@
import gym
import pytest
from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
from torchy_baselines.common.vec_env import DummyVecEnv
MODEL_LIST = [
CEMRL,
PPO,
A2C,
TD3,
SAC,
]
@pytest.mark.parametrize("model_class", MODEL_LIST)
def test_auto_wrap(model_class):
# test auto wrapping of env into a VecEnv
env = gym.make('Pendulum-v0')
eval_env = gym.make('Pendulum-v0')
model = model_class('MlpPolicy', env)
model.learn(100, eval_env=eval_env)
@pytest.mark.parametrize("model_class", MODEL_LIST)
def test_predict(model_class):
# test detection of different shapes by the predict method
model = model_class('MlpPolicy', 'Pendulum-v0')
env = gym.make('Pendulum-v0')
vec_env = DummyVecEnv([lambda: gym.make('Pendulum-v0'), lambda: gym.make('Pendulum-v0')])
obs = env.reset()
action = model.predict(obs)
assert action.shape == env.action_space.shape
assert env.action_space.contains(action)
vec_env_obs = vec_env.reset()
action = model.predict(vec_env_obs)
assert action.shape[0] == vec_env_obs.shape[0]
@pytest.mark.parametrize("model_class", [A2C, PPO])
def test_predict_discrete(model_class):
# test detection of different shapes by the predict method
model = model_class('MlpPolicy', 'CartPole-v1')
env = gym.make('CartPole-v1')
vec_env = DummyVecEnv([lambda: gym.make('CartPole-v1'), lambda: gym.make('CartPole-v1')])
obs = env.reset()
action = model.predict(obs)
assert action.shape == ()
assert env.action_space.contains(action)
vec_env_obs = vec_env.reset()
action = model.predict(vec_env_obs)
assert action.shape[0] == vec_env_obs.shape[0]

12
tests/test_run.py
View file

@ -14,18 +14,12 @@ def test_td3(action_noise):
model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]),
learning_starts=100, verbose=1, create_eval_env=True, action_noise=action_noise)
model.learn(total_timesteps=1000, eval_freq=500)
model.save("test_save")
model.load("test_save")
os.remove("test_save.zip")
def test_cemrl():
model = CEMRL('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), pop_size=2, n_grad=1,
learning_starts=100, verbose=1, create_eval_env=True, action_noise=action_noise)
model.learn(total_timesteps=1000, eval_freq=500)
model.save("test_save")
model.load("test_save")
os.remove("test_save.zip")
@pytest.mark.parametrize("model_class", [A2C, PPO])
@ -33,9 +27,6 @@ def test_cemrl():
def test_onpolicy(model_class, env_id):
model = model_class('MlpPolicy', env_id, policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True)
model.learn(total_timesteps=1000, eval_freq=500)
model.save("test_save")
model.load("test_save")
os.remove("test_save.zip")
@pytest.mark.parametrize("ent_coef", ['auto', 0.01])
@ -44,6 +35,3 @@ def test_sac(ent_coef):
learning_starts=100, verbose=1, create_eval_env=True, ent_coef=ent_coef,
action_noise=NormalActionNoise(np.zeros(1), np.zeros(1)))
model.learn(total_timesteps=1000, eval_freq=500)
model.save("test_save")
model.load("test_save")
os.remove("test_save.zip")

1
tests/test_save_load.py
View file

@ -134,6 +134,7 @@ def test_exclude_include_saved_params(model_class):
# clear file from os
os.remove("test_save.zip")
@pytest.mark.parametrize("model_class", [SAC, TD3])
def test_save_load_replay_buffer(model_class):
log_folder = 'logs'

37
torchy_baselines/common/base_class.py
View file

@ -335,24 +335,25 @@ class BaseRLModel(ABC):
else:
raise ValueError("Error: Unexpected observation shape {} for ".format(observation.shape) +
"Discrete environment, please use (1,) or (n_env, 1) for the observation shape.")
elif isinstance(observation_space, gym.spaces.MultiDiscrete):
if observation.shape == (len(observation_space.nvec),):
return False
elif len(observation.shape) == 2 and observation.shape[1] == len(observation_space.nvec):
return True
else:
raise ValueError("Error: Unexpected observation shape {} for MultiDiscrete ".format(observation.shape) +
"environment, please use ({},) or ".format(len(observation_space.nvec)) +
"(n_env, {}) for the observation shape.".format(len(observation_space.nvec)))
elif isinstance(observation_space, gym.spaces.MultiBinary):
if observation.shape == (observation_space.n,):
return False
elif len(observation.shape) == 2 and observation.shape[1] == observation_space.n:
return True
else:
raise ValueError("Error: Unexpected observation shape {} for MultiBinary ".format(observation.shape) +
"environment, please use ({},) or ".format(observation_space.n) +
"(n_env, {}) for the observation shape.".format(observation_space.n))
# TODO: add support for MultiDiscrete and MultiBinary action spaces
# elif isinstance(observation_space, gym.spaces.MultiDiscrete):
# if observation.shape == (len(observation_space.nvec),):
# return False
# elif len(observation.shape) == 2 and observation.shape[1] == len(observation_space.nvec):
# return True
# else:
# raise ValueError("Error: Unexpected observation shape {} for MultiDiscrete ".format(observation.shape) +
# "environment, please use ({},) or ".format(len(observation_space.nvec)) +
# "(n_env, {}) for the observation shape.".format(len(observation_space.nvec)))
# elif isinstance(observation_space, gym.spaces.MultiBinary):
# if observation.shape == (observation_space.n,):
# return False
# elif len(observation.shape) == 2 and observation.shape[1] == observation_space.n:
# return True
# else:
# raise ValueError("Error: Unexpected observation shape {} for MultiBinary ".format(observation.shape) +
# "environment, please use ({},) or ".format(observation_space.n) +
# "(n_env, {}) for the observation shape.".format(observation_space.n))
else:
raise ValueError("Error: Cannot determine if the observation is vectorized with the space type {}."
.format(observation_space))