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)
def test_channel_first_env(tmp_path):
# test_cnn uses environment with HxWxC setup that is transposed, but we
# also want to work with CxHxW envs directly without transposing wrapper.
SAVE_NAME = "cnn_model.zip"
# Create environment with transposed images (CxHxW).
# If underlying CNN processes the data in wrong format,
# it will raise an error of negative dimension sizes while creating convolutions
env = FakeImageEnv(screen_height=40,
screen_width=40,
n_channels=1,
discrete=True,
channel_first=True)
model = A2C("CnnPolicy", env, n_steps=100).learn(250)
assert not is_vecenv_wrapped(model.get_env(), VecTransposeImage)
obs = env.reset()
action, _ = model.predict(obs, deterministic=True)
model.save(tmp_path / SAVE_NAME)
del model
model = A2C.load(tmp_path / SAVE_NAME)
# Check that the prediction is the same
assert np.allclose(action, model.predict(obs, deterministic=True)[0])
os.remove(str(tmp_path / SAVE_NAME))
def test_cnn(tmp_path, model_class):
SAVE_NAME = "cnn_model.zip"
# 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=64)
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)
# FakeImageEnv is channel last by default and should be wrapped
assert is_vecenv_wrapped(model.get_env(), VecTransposeImage)
obs = env.reset()
action, _ = model.predict(obs, deterministic=True)
model.save(tmp_path / SAVE_NAME)
del model
model = model_class.load(tmp_path / SAVE_NAME)
# Check that the prediction is the same
assert np.allclose(action, model.predict(obs, deterministic=True)[0])
os.remove(str(tmp_path / SAVE_NAME))
def test_cnn(tmp_path, model_class):
SAVE_NAME = "cnn_model.zip"
# 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 {TQC})
kwargs = {}
if model_class in {TQC, QRDQN}:
# Avoid memory error when using replay buffer
# Reduce the size of the features and the number of quantiles
kwargs = dict(
buffer_size=250,
policy_kwargs=dict(
n_quantiles=25,
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(tmp_path / SAVE_NAME)
del model
model = model_class.load(tmp_path / SAVE_NAME)
# Check that the prediction is the same
assert np.allclose(action, model.predict(obs, deterministic=True)[0])
os.remove(str(tmp_path / SAVE_NAME))
def test_cnn(tmp_path, model_class):
SAVE_NAME = "cnn_model.zip"
# 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 {TQC},
)
kwargs = {}
if model_class in {TQC, QRDQN}:
# Avoid memory error when using replay buffer
# Reduce the size of the features and the number of quantiles
kwargs = dict(
buffer_size=250,
policy_kwargs=dict(
n_quantiles=25,
features_extractor_kwargs=dict(features_dim=32),
),
)
else:
kwargs = dict(
buffer_size=250,
policy_kwargs=dict(features_extractor_kwargs=dict(
features_dim=32)),
seed=1,
)
model = model_class("CnnPolicy", env, **kwargs).learn(250)
obs = env.reset()
# FakeImageEnv is channel last by default and should be wrapped
assert is_vecenv_wrapped(model.get_env(), VecTransposeImage)
# Test stochastic predict with channel last input
if model_class in {QRDQN, DQNClipped, DQNReg}:
model.exploration_rate = 0.9
for _ in range(10):
model.predict(obs, deterministic=False)
action, _ = model.predict(obs, deterministic=True)
model.save(tmp_path / SAVE_NAME)
del model
model = model_class.load(tmp_path / SAVE_NAME)
# Check that the prediction is the same
assert np.allclose(action, model.predict(obs, deterministic=True)[0])
os.remove(str(tmp_path / SAVE_NAME))
def test_non_default_spaces(new_obs_space):
env = FakeImageEnv()
env.observation_space = new_obs_space
# Patch methods to avoid errors
env.reset = new_obs_space.sample
def patched_step(_action):
return new_obs_space.sample(), 0.0, False, {}
env.step = patched_step
with pytest.warns(UserWarning):
check_env(env)
def test_high_dimension_action_space():
"""
Test for continuous action space
with more than one action.
"""
env = FakeImageEnv()
# Patch the action space
env.action_space = spaces.Box(low=-1, high=1, shape=(20,), dtype=np.float32)
# Patch to avoid error
def patched_step(_action):
return env.observation_space.sample(), 0.0, False, {}
env.step = patched_step
check_env(env)
def test_save_load_env_cnn(tmp_path, model_class):
"""
Test loading with an env that requires a ``CnnPolicy``.
This is to test wrapping and observation space check.
We test one on-policy and one off-policy
algorithm as the rest share the loading part.
"""
env = FakeImageEnv(screen_height=40,
screen_width=40,
n_channels=2,
discrete=False)
kwargs = dict(policy_kwargs=dict(net_arch=[32]))
if model_class == TD3:
kwargs.update(dict(buffer_size=100, learning_starts=50))
model = model_class("CnnPolicy", env, **kwargs).learn(100)
model.save(tmp_path / "test_save")
# Test loading with env and continuing training
model = model_class.load(str(tmp_path / "test_save.zip"),
env=env).learn(100)
# clear file from os
os.remove(tmp_path / "test_save.zip")
def test_save_load_policy(tmp_path, model_class, policy_str):
"""
Test saving and loading policy only.
:param model_class: (BaseAlgorithm) A RL model
:param policy_str: (str) Name of the policy.
"""
kwargs = {}
if policy_str == "MlpPolicy":
env = select_env(model_class)
else:
if model_class in [SAC, TD3, DQN]:
# Avoid memory error when using replay buffer
# Reduce the size of the features
kwargs = dict(buffer_size=250)
env = FakeImageEnv(screen_height=40, screen_width=40, n_channels=2, discrete=model_class == DQN)
env = DummyVecEnv([lambda: env])
# create model
model = model_class(policy_str, env, policy_kwargs=dict(net_arch=[16]), verbose=1, **kwargs)
model.learn(total_timesteps=500)
env.reset()
observations = np.concatenate([env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0)
policy = model.policy
policy_class = policy.__class__
actor, actor_class = None, None
if model_class in [SAC, TD3]:
actor = policy.actor
actor_class = actor.__class__
# Get dictionary of current parameters
params = deepcopy(policy.state_dict())
# Modify all parameters to be random values
random_params = dict((param_name, th.rand_like(param)) for param_name, param in params.items())
# Update model parameters with the new random values
policy.load_state_dict(random_params)
new_params = policy.state_dict()
# Check that all params are different now
for k in params:
assert not th.allclose(params[k], new_params[k]), "Parameters did not change as expected."
params = new_params
# get selected actions
selected_actions, _ = policy.predict(observations, deterministic=True)
# Should also work with the actor only
if actor is not None:
selected_actions_actor, _ = actor.predict(observations, deterministic=True)
# Save and load policy
policy.save(tmp_path / "policy.pkl")
# Save and load actor
if actor is not None:
actor.save(tmp_path / "actor.pkl")
del policy, actor
policy = policy_class.load(tmp_path / "policy.pkl")
if actor_class is not None:
actor = actor_class.load(tmp_path / "actor.pkl")
# check if params are still the same after load
new_params = policy.state_dict()
# Check that all params are the same as before save load procedure now
for key in params:
assert th.allclose(params[key], new_params[key]), "Policy parameters not the same after save and load."
# check if model still selects the same actions
new_selected_actions, _ = policy.predict(observations, deterministic=True)
assert np.allclose(selected_actions, new_selected_actions, 1e-4)
if actor_class is not None:
new_selected_actions_actor, _ = actor.predict(observations, deterministic=True)
assert np.allclose(selected_actions_actor, new_selected_actions_actor, 1e-4)
assert np.allclose(selected_actions_actor, new_selected_actions, 1e-4)
# clear file from os
os.remove(tmp_path / "policy.pkl")
if actor_class is not None:
os.remove(tmp_path / "actor.pkl")
def test_save_load_q_net(tmp_path, model_class, policy_str):
"""
Test saving and loading q-network/quantile net only.
:param model_class: (BaseAlgorithm) A RL model
:param policy_str: (str) Name of the policy.
"""
kwargs = dict(policy_kwargs=dict(net_arch=[16]))
if policy_str == "MlpPolicy":
env = select_env(model_class)
else:
if model_class in [DQN]:
# Avoid memory error when using replay buffer
# Reduce the size of the features
kwargs = dict(
buffer_size=250,
learning_starts=100,
policy_kwargs=dict(features_extractor_kwargs=dict(
features_dim=32)),
)
env = FakeImageEnv(screen_height=40,
screen_width=40,
n_channels=2,
discrete=model_class == DQN)
env = DummyVecEnv([lambda: env])
# create model
model = model_class(policy_str, env, verbose=1, **kwargs)
model.learn(total_timesteps=300)
env.reset()
observations = np.concatenate(
[env.step([env.action_space.sample()])[0] for _ in range(10)], axis=0)
q_net = model.q_net
q_net_class = q_net.__class__
# Get dictionary of current parameters
params = deepcopy(q_net.state_dict())
# Modify all parameters to be random values
random_params = dict((param_name, th.rand_like(param))
for param_name, param in params.items())
# Update model parameters with the new random values
q_net.load_state_dict(random_params)
new_params = q_net.state_dict()
# Check that all params are different now
for k in params:
assert not th.allclose(
params[k], new_params[k]), "Parameters did not change as expected."
params = new_params
# get selected actions
selected_actions, _ = q_net.predict(observations, deterministic=True)
# Save and load q_net
q_net.save(tmp_path / "q_net.pkl")
del q_net
q_net = q_net_class.load(tmp_path / "q_net.pkl")
# check if params are still the same after load
new_params = q_net.state_dict()
# Check that all params are the same as before save load procedure now
for key in params:
assert th.allclose(
params[key], new_params[key]
), "Policy parameters not the same after save and load."
# check if model still selects the same actions
new_selected_actions, _ = q_net.predict(observations, deterministic=True)
assert np.allclose(selected_actions, new_selected_actions, 1e-4)
# clear file from os
os.remove(tmp_path / "q_net.pkl")
def test_features_extractor_target_net(model_class, share_features_extractor):
if model_class == DQN and share_features_extractor:
pytest.skip()
env = FakeImageEnv(screen_height=40,
screen_width=40,
n_channels=1,
discrete=model_class not in {SAC, TD3})
# Avoid memory error when using replay buffer
# Reduce the size of the features
kwargs = dict(
buffer_size=250,
learning_starts=100,
policy_kwargs=dict(features_extractor_kwargs=dict(features_dim=32)))
if model_class != DQN:
kwargs["policy_kwargs"][
"share_features_extractor"] = share_features_extractor
# No delay for TD3 (changes when the actor and polyak update take place)
if model_class == TD3:
kwargs["policy_delay"] = 1
model = model_class("CnnPolicy", env, seed=0, **kwargs)
patch_dqn_names_(model)
if share_features_extractor:
# Check that the objects are the same and not just copied
assert id(model.policy.actor.features_extractor) == id(
model.policy.critic.features_extractor)
if model_class == TD3:
assert id(model.policy.actor_target.features_extractor) == id(
model.policy.critic_target.features_extractor)
# Actor and critic feature extractor should be the same
td3_features_extractor_check = check_td3_feature_extractor_match
else:
# Actor and critic feature extractor should differ same
td3_features_extractor_check = check_td3_feature_extractor_differ
# Check that the object differ
if model_class != DQN:
assert id(model.policy.actor.features_extractor) != id(
model.policy.critic.features_extractor)
if model_class == TD3:
assert id(model.policy.actor_target.features_extractor) != id(
model.policy.critic_target.features_extractor)
# Critic and target should be equal at the begginning of training
params_should_match(model.critic.parameters(),
model.critic_target.parameters())
# TD3 has also a target actor net
if model_class == TD3:
params_should_match(model.actor.parameters(),
model.actor_target.parameters())
model.learn(200)
# Critic and target should differ
params_should_differ(model.critic.parameters(),
model.critic_target.parameters())
if model_class == TD3:
params_should_differ(model.actor.parameters(),
model.actor_target.parameters())
td3_features_extractor_check(model)
# Re-initialize and collect some random data (without doing gradient steps,
# since 10 < learning_starts = 100)
model = model_class("CnnPolicy", env, seed=0, **kwargs).learn(10)
patch_dqn_names_(model)
original_param = deepcopy(list(model.critic.parameters()))
original_target_param = deepcopy(list(model.critic_target.parameters()))
if model_class == TD3:
original_actor_target_param = deepcopy(
list(model.actor_target.parameters()))
# Deactivate copy to target
model.tau = 0.0
model.train(gradient_steps=1)
# Target should be the same
params_should_match(original_target_param,
model.critic_target.parameters())
if model_class == TD3:
params_should_match(original_actor_target_param,
model.actor_target.parameters())
td3_features_extractor_check(model)
# not the same for critic net (updated by gradient descent)
params_should_differ(original_param, model.critic.parameters())
# Update the reference as it should not change in the next step
original_param = deepcopy(list(model.critic.parameters()))
if model_class == TD3:
original_actor_param = deepcopy(list(model.actor.parameters()))
# Deactivate learning rate
model.lr_schedule = lambda _: 0.0
# Re-activate polyak update
model.tau = 0.01
# Special case for DQN: target net is updated in the `collect_rollouts()`
# not the `train()` method
if model_class == DQN:
model.target_update_interval = 1
model._on_step()
model.train(gradient_steps=1)
# Target should have changed now (due to polyak update)
params_should_differ(original_target_param,
model.critic_target.parameters())
# Critic should be the same
params_should_match(original_param, model.critic.parameters())
if model_class == TD3:
params_should_differ(original_actor_target_param,
model.actor_target.parameters())
params_should_match(original_actor_param, model.actor.parameters())
td3_features_extractor_check(model)
def test_feature_extractor_target_net(model_class, share_features_extractor):
if (model_class == QRDQN or model_class == DQNReg
or model_class == DQNClipped) and share_features_extractor:
pytest.skip()
env = FakeImageEnv(screen_height=40,
screen_width=40,
n_channels=1,
discrete=model_class not in {TQC})
if model_class in {TQC, QRDQN}:
# Avoid memory error when using replay buffer
# Reduce the size of the features and the number of quantiles
kwargs = dict(
buffer_size=250,
learning_starts=100,
policy_kwargs=dict(
n_quantiles=25,
features_extractor_kwargs=dict(features_dim=32)),
)
else:
kwargs = dict(buffer_size=250,
learning_starts=100,
policy_kwargs=dict(features_extractor_kwargs=dict(
features_dim=32)))
if model_class not in {QRDQN, DQNClipped, DQNReg}:
kwargs["policy_kwargs"][
"share_features_extractor"] = share_features_extractor
model = model_class("CnnPolicy", env, seed=0, **kwargs)
patch_qrdqn_names_(model)
if share_features_extractor:
# Check that the objects are the same and not just copied
assert id(model.policy.actor.features_extractor) == id(
model.policy.critic.features_extractor)
else:
# Check that the objects differ
if model_class not in {QRDQN, DQNClipped, DQNReg}:
assert id(model.policy.actor.features_extractor) != id(
model.policy.critic.features_extractor)
# Critic and target should be equal at the beginning of training
params_should_match(model.critic.parameters(),
model.critic_target.parameters())
model.learn(200)
# Critic and target should differ
params_should_differ(model.critic.parameters(),
model.critic_target.parameters())
# Re-initialize and collect some random data (without doing gradient steps)
model = model_class("CnnPolicy", env, seed=0, **kwargs).learn(10)
patch_qrdqn_names_(model)
original_param = deepcopy(list(model.critic.parameters()))
original_target_param = deepcopy(list(model.critic_target.parameters()))
# Deactivate copy to target
model.tau = 0.0
model.train(gradient_steps=1)
# Target should be the same
params_should_match(original_target_param,
model.critic_target.parameters())
# not the same for critic net (updated by gradient descent)
params_should_differ(original_param, model.critic.parameters())
# Update the reference as it should not change in the next step
original_param = deepcopy(list(model.critic.parameters()))
# Deactivate learning rate
model.lr_schedule = lambda _: 0.0
# Re-activate polyak update
model.tau = 0.01
# Special case for QRDQN: target net is updated in the `collect_rollouts()`
# not the `train()` method
if model_class in {QRDQN, DQNClipped, DQNReg}:
model.target_update_interval = 1
model._on_step()
model.train(gradient_steps=1)
# Target should have changed now (due to polyak update)
params_should_differ(original_target_param,
model.critic_target.parameters())
# Critic should be the same
params_should_match(original_param, model.critic.parameters())