def test_time_feature(): env = gym.make("Pendulum-v0") env = TimeFeatureWrapper(env) check_env(env, warn=False) # Check for four episodes max_timesteps = 200 obs = env.reset() for _ in range(4): check_time_feature(obs, timestep=0, max_timesteps=max_timesteps) for step in range(1, max_timesteps + 1): obs, _, done, _ = env.step(env.action_space.sample()) check_time_feature(obs, timestep=step, max_timesteps=max_timesteps) if done: obs = env.reset() env = BitFlippingEnv() with pytest.raises(AssertionError): env = TimeFeatureWrapper(env) env = CustomGoalEnv() env = TimeFeatureWrapper(env, max_steps=500) obs = env.reset() check_time_feature(obs["observation"], timestep=0, max_timesteps=500) obs, _, _, _ = env.step(env.action_space.sample()) check_time_feature(obs["observation"], timestep=1, max_timesteps=500) # In test mode, the time feature must be constant env = gym.make("Pendulum-v0") env = TimeFeatureWrapper(env, test_mode=True) obs = env.reset() check_time_feature(obs, timestep=0, max_timesteps=200) obs, _, _, _ = env.step(env.action_space.sample()) # Should be the same check_time_feature(obs, timestep=0, max_timesteps=200)
def test_vec_normalize(model_class): """ Additional tests for PPO/A2C/SAC/DDPG/TD3/DQN to check observation space support for GoalEnv and VecNormalize using MultiInputPolicy. """ env = DummyVecEnv([ lambda: BitFlippingEnv(n_bits=4, continuous=not (model_class == DQN)) ]) env = VecNormalize(env) kwargs = {} n_steps = 256 if model_class in {A2C, PPO}: kwargs = dict( n_steps=128, policy_kwargs=dict(net_arch=[32], ), ) else: # Avoid memory error when using replay buffer # Reduce the size of the features and make learning faster kwargs = dict( buffer_size=250, policy_kwargs=dict(net_arch=[32], ), train_freq=8, gradient_steps=1, ) if model_class == DQN: kwargs["learning_starts"] = 0 model = model_class("MultiInputPolicy", env, gamma=0.5, seed=1, **kwargs) model.learn(total_timesteps=n_steps) evaluate_policy(model, env, n_eval_episodes=5, warn=False)
def test_performance_her(online_sampling, n_bits): """ That DQN+HER can solve BitFlippingEnv. It should not work when n_sampled_goal=0 (DQN alone). """ env = BitFlippingEnv(n_bits=n_bits, continuous=False) model = DQN( "MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict( n_sampled_goal=5, goal_selection_strategy="future", online_sampling=online_sampling, max_episode_length=n_bits, ), verbose=1, learning_rate=5e-4, train_freq=1, learning_starts=100, exploration_final_eps=0.02, target_update_interval=500, seed=0, batch_size=32, buffer_size=int(1e5), ) model.learn(total_timesteps=5000, log_interval=50) # 90% training success assert np.mean(model.ep_success_buffer) > 0.90
def test_goal_selection_strategy(goal_selection_strategy, online_sampling): """ Test different goal strategies. """ env = BitFlippingEnv(continuous=True) normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)) model = SAC( "MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict( goal_selection_strategy=goal_selection_strategy, online_sampling=online_sampling, max_episode_length=10, n_sampled_goal=2, ), train_freq=4, gradient_steps=1, policy_kwargs=dict(net_arch=[64]), learning_starts=100, buffer_size=int(1e5), action_noise=normal_action_noise, ) assert model.action_noise is not None model.learn(total_timesteps=150)
def test_full_replay_buffer(): """ Test if HER works correctly with a full replay buffer when using online sampling. It should not sample the current episode which is not finished. """ n_bits = 4 env = BitFlippingEnv(n_bits=n_bits, continuous=True) # use small buffer size to get the buffer full model = SAC( "MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict( n_sampled_goal=2, goal_selection_strategy="future", online_sampling=True, max_episode_length=n_bits, ), gradient_steps=1, train_freq=4, policy_kwargs=dict(net_arch=[64]), learning_starts=1, buffer_size=20, verbose=1, seed=757, ) model.learn(total_timesteps=100)
def test_her(model_class, online_sampling, image_obs_space): """ Test Hindsight Experience Replay. """ n_bits = 4 env = BitFlippingEnv( n_bits=n_bits, continuous=not (model_class == DQN), image_obs_space=image_obs_space, ) model = model_class( "MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict( n_sampled_goal=2, goal_selection_strategy="future", online_sampling=online_sampling, max_episode_length=n_bits, ), train_freq=4, gradient_steps=1, policy_kwargs=dict(net_arch=[64]), learning_starts=100, buffer_size=int(2e4), ) model.learn(total_timesteps=150) evaluate_policy(model, Monitor(env))
def test_bit_flipping(kwargs): # Additional tests for BitFlippingEnv env = BitFlippingEnv(**kwargs) with pytest.warns(None) as record: check_env(env) # No warnings for custom envs assert len(record) == 0
def test_bit_flipping(kwargs): # Additional tests for BitFlippingEnv env = BitFlippingEnv(**kwargs) with warnings.catch_warnings(record=True) as record: check_env(env) # No warnings for custom envs assert len(record) == 0
def test_get_max_episode_length(): dict_env = DummyVecEnv([lambda: BitFlippingEnv()]) # Cannot infer max epsiode length with pytest.raises(ValueError): get_time_limit(dict_env, current_max_episode_length=None) default_length = 10 assert get_time_limit( dict_env, current_max_episode_length=default_length) == default_length env = gym.make("CartPole-v1") vec_env = DummyVecEnv([lambda: env]) assert get_time_limit(vec_env, current_max_episode_length=None) == 500 # Overwrite max_episode_steps assert get_time_limit( vec_env, current_max_episode_length=default_length) == default_length # Set max_episode_steps to None env.spec.max_episode_steps = None vec_env = DummyVecEnv([lambda: env]) with pytest.raises(ValueError): get_time_limit(vec_env, current_max_episode_length=None) # Initialize HER and specify max_episode_length, should not raise an issue DQN("MultiInputPolicy", dict_env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict(max_episode_length=5)) with pytest.raises(ValueError): DQN("MultiInputPolicy", dict_env, replay_buffer_class=HerReplayBuffer) # Wrapped in a timelimit, should be fine # Note: it requires env.spec to be defined env = DummyVecEnv([lambda: gym.wrappers.TimeLimit(BitFlippingEnv(), 10)]) DQN("MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict(max_episode_length=5))
def test_eval_success_logging(tmp_path): n_bits = 2 env = BitFlippingEnv(n_bits=n_bits) eval_env = DummyVecEnv([lambda: BitFlippingEnv(n_bits=n_bits)]) eval_callback = EvalCallback( eval_env, eval_freq=250, log_path=tmp_path, warn=False, ) model = DQN( "MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, learning_starts=100, seed=0, replay_buffer_kwargs=dict(max_episode_length=n_bits), ) model.learn(500, callback=eval_callback) assert len(eval_callback._is_success_buffer) > 0 # More than 50% success rate assert np.mean(eval_callback._is_success_buffer) > 0.5
def test_goal_env(model_class): env = BitFlippingEnv(n_bits=4) # check that goal env works for PPO/A2C that cannot use HER replay buffer model = model_class("MultiInputPolicy", env, n_steps=64).learn(250) evaluate_policy(model, model.get_env())
def test_policy_hint(policy): # Common mistake: using the wrong policy with pytest.raises(ValueError): PPO(policy, BitFlippingEnv(n_bits=4))
def test_save_load_replay_buffer(tmp_path, recwarn, online_sampling, truncate_last_trajectory): """ Test if 'save_replay_buffer' and 'load_replay_buffer' works correctly """ # remove gym warnings warnings.filterwarnings(action="ignore", category=DeprecationWarning) warnings.filterwarnings(action="ignore", category=UserWarning, module="gym") path = pathlib.Path(tmp_path / "replay_buffer.pkl") path.parent.mkdir(exist_ok=True, parents=True) # to not raise a warning env = BitFlippingEnv(n_bits=4, continuous=True) model = SAC( "MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict( n_sampled_goal=2, goal_selection_strategy="future", online_sampling=online_sampling, max_episode_length=4, ), gradient_steps=1, train_freq=4, buffer_size=int(2e4), policy_kwargs=dict(net_arch=[64]), seed=1, ) model.learn(200) if online_sampling: old_replay_buffer = deepcopy(model.replay_buffer) else: old_replay_buffer = deepcopy(model.replay_buffer.replay_buffer) model.save_replay_buffer(path) del model.replay_buffer with pytest.raises(AttributeError): model.replay_buffer # Check that there is no warning assert len(recwarn) == 0 model.load_replay_buffer(path, truncate_last_traj=truncate_last_trajectory) if truncate_last_trajectory: assert len(recwarn) == 1 warning = recwarn.pop(UserWarning) assert "The last trajectory in the replay buffer will be truncated" in str( warning.message) else: assert len(recwarn) == 0 if online_sampling: n_episodes_stored = model.replay_buffer.n_episodes_stored assert np.allclose( old_replay_buffer._buffer["observation"][:n_episodes_stored], model.replay_buffer._buffer["observation"][:n_episodes_stored], ) assert np.allclose( old_replay_buffer._buffer["next_obs"][:n_episodes_stored], model.replay_buffer._buffer["next_obs"][:n_episodes_stored], ) assert np.allclose( old_replay_buffer._buffer["action"][:n_episodes_stored], model.replay_buffer._buffer["action"][:n_episodes_stored], ) assert np.allclose( old_replay_buffer._buffer["reward"][:n_episodes_stored], model.replay_buffer._buffer["reward"][:n_episodes_stored], ) # we might change the last done of the last trajectory so we don't compare it assert np.allclose( old_replay_buffer._buffer["done"][:n_episodes_stored - 1], model.replay_buffer._buffer["done"][:n_episodes_stored - 1], ) else: replay_buffer = model.replay_buffer.replay_buffer assert np.allclose(old_replay_buffer.observations["observation"], replay_buffer.observations["observation"]) assert np.allclose(old_replay_buffer.observations["desired_goal"], replay_buffer.observations["desired_goal"]) assert np.allclose(old_replay_buffer.actions, replay_buffer.actions) assert np.allclose(old_replay_buffer.rewards, replay_buffer.rewards) assert np.allclose(old_replay_buffer.dones, replay_buffer.dones) # test if continuing training works properly reset_num_timesteps = False if truncate_last_trajectory is False else True model.learn(200, reset_num_timesteps=reset_num_timesteps)
def test_save_load(tmp_path, model_class, use_sde, online_sampling): """ Test if 'save' and 'load' saves and loads model correctly """ if use_sde and model_class != SAC: pytest.skip("Only SAC has gSDE support") n_bits = 4 env = BitFlippingEnv(n_bits=n_bits, continuous=not (model_class == DQN)) kwargs = dict(use_sde=True) if use_sde else {} # create model model = model_class("MultiInputPolicy", env, replay_buffer_class=HerReplayBuffer, replay_buffer_kwargs=dict( n_sampled_goal=2, goal_selection_strategy="future", online_sampling=online_sampling, max_episode_length=n_bits, ), verbose=0, tau=0.05, batch_size=128, learning_rate=0.001, policy_kwargs=dict(net_arch=[64]), buffer_size=int(1e5), gamma=0.98, gradient_steps=1, train_freq=4, learning_starts=100, **kwargs) model.learn(total_timesteps=150) obs = env.reset() observations = {key: [] for key in obs.keys()} for _ in range(10): obs = env.step(env.action_space.sample())[0] for key in obs.keys(): observations[key].append(obs[key]) observations = {key: np.array(obs) for key, obs in observations.items()} # Get dictionary of current parameters params = deepcopy(model.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 model.policy.load_state_dict(random_params) new_params = model.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, _ = model.predict(observations, deterministic=True) # Check model.save(tmp_path / "test_save.zip") del model # test custom_objects # Load with custom objects custom_objects = dict(learning_rate=2e-5, dummy=1.0) model_ = model_class.load(str(tmp_path / "test_save.zip"), env=env, custom_objects=custom_objects, verbose=2) assert model_.verbose == 2 # Check that the custom object was taken into account assert model_.learning_rate == custom_objects["learning_rate"] # Check that only parameters that are here already are replaced assert not hasattr(model_, "dummy") model = model_class.load(str(tmp_path / "test_save.zip"), env=env) # check if params are still the same after load new_params = model.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] ), "Model parameters not the same after save and load." # check if model still selects the same actions new_selected_actions, _ = model.predict(observations, deterministic=True) assert np.allclose(selected_actions, new_selected_actions, 1e-4) # check if learn still works model.learn(total_timesteps=150) # Test that the change of parameters works model = model_class.load(str(tmp_path / "test_save.zip"), env=env, verbose=3, learning_rate=2.0) assert model.learning_rate == 2.0 assert model.verbose == 3 # clear file from os os.remove(tmp_path / "test_save.zip")