diff --git a/tests/test_predict.py b/tests/test_predict.py new file mode 100644 index 0000000..e75deec --- /dev/null +++ b/tests/test_predict.py @@ -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] diff --git a/tests/test_run.py b/tests/test_run.py index 1d206c9..fdfcff6 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -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") diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 45c0ac5..9d73ddf 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -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' diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index eabc7f1..15da3bb 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -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))