mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-05-18 21:30:19 +00:00
b702884c23
* Modified actor-critic policies & MlpExtractor class ActorCriticPolicy: - changed type hint of net_arch param: now it's a dict - removed check that if features extractor is not shared: no shared layers are allowed in the mlp_extractor regardless of the features extractor ActorCriticCnnPolicy: - changed type hint of net_arch param: now it's a dict MultiInputActorcriticPolicy: - changed type hint of net_arch param: now it's a dict MlpExtractor: - changed type hint of net_arch param: now it's a dict - adapted networks creation - adapted methods: forward, forward_actor & forward_critic * Removed shared layers in mlp_extractor * Updated docs and changelog + reformat * Updated custom policy tests * Removed test on deprecation warning for share layers in mlp_extractor Now shared layers are removed * Update version * Update RL Zoo doc * Fix linter warnings * Add ruff to Makefile (experimental) * Add backward compat code and minor updates * Update tests * Add backward compatibility * Fix test * Improve compat code Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>
53 lines
2 KiB
Python
53 lines
2 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from stable_baselines3 import A2C, DDPG, DQN, PPO, SAC, TD3
|
|
from stable_baselines3.common.envs import IdentityEnv, IdentityEnvBox, IdentityEnvMultiBinary, IdentityEnvMultiDiscrete
|
|
from stable_baselines3.common.evaluation import evaluate_policy
|
|
from stable_baselines3.common.noise import NormalActionNoise
|
|
from stable_baselines3.common.vec_env import DummyVecEnv
|
|
|
|
DIM = 4
|
|
|
|
|
|
@pytest.mark.parametrize("model_class", [A2C, PPO, DQN])
|
|
@pytest.mark.parametrize("env", [IdentityEnv(DIM), IdentityEnvMultiDiscrete(DIM), IdentityEnvMultiBinary(DIM)])
|
|
def test_discrete(model_class, env):
|
|
env_ = DummyVecEnv([lambda: env])
|
|
kwargs = {}
|
|
n_steps = 3000
|
|
if model_class == DQN:
|
|
kwargs = dict(learning_starts=0)
|
|
n_steps = 4000
|
|
# DQN only support discrete actions
|
|
if isinstance(env, (IdentityEnvMultiDiscrete, IdentityEnvMultiBinary)):
|
|
return
|
|
elif model_class == A2C:
|
|
# slightly higher budget
|
|
n_steps = 3500
|
|
|
|
model = model_class("MlpPolicy", env_, gamma=0.4, seed=1, **kwargs).learn(n_steps)
|
|
|
|
evaluate_policy(model, env_, n_eval_episodes=20, reward_threshold=90, warn=False)
|
|
obs = env.reset()
|
|
|
|
assert np.shape(model.predict(obs)[0]) == np.shape(obs)
|
|
|
|
|
|
@pytest.mark.parametrize("model_class", [A2C, PPO, SAC, DDPG, TD3])
|
|
def test_continuous(model_class):
|
|
env = IdentityEnvBox(eps=0.5)
|
|
|
|
n_steps = {A2C: 3500, PPO: 3000, SAC: 700, TD3: 500, DDPG: 500}[model_class]
|
|
|
|
kwargs = dict(policy_kwargs=dict(net_arch=[64, 64]), seed=0, gamma=0.95)
|
|
if model_class in [TD3]:
|
|
n_actions = 1
|
|
action_noise = NormalActionNoise(mean=np.zeros(n_actions), sigma=0.1 * np.ones(n_actions))
|
|
kwargs["action_noise"] = action_noise
|
|
elif model_class in [A2C]:
|
|
kwargs["policy_kwargs"]["log_std_init"] = -0.5
|
|
|
|
model = model_class("MlpPolicy", env, **kwargs).learn(n_steps)
|
|
|
|
evaluate_policy(model, env, n_eval_episodes=20, reward_threshold=90, warn=False)
|