mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-07-07 04:39:24 +00:00
Merge pull request #6 from Antonin-Raffin/feat/sde-features
Feature Extract for SDE
This commit is contained in:
Raffin, Antonin
GitHub Enterprise
commit
358b27e9c9
30 changed files with 1102 additions and 383 deletions
21
README.md
21
README.md
|
|
@ -14,24 +14,9 @@ PyTorch version of [Stable Baselines](https://github.com/hill-a/stable-baselines
|
|||
- SAC
|
||||
- TD3
|
||||
|
||||
- SDE support for A2C, PPO, SAC and TD3.
|
||||
|
||||
|
||||
## Roadmap
|
||||
|
||||
TODO:
|
||||
- save/load
|
||||
- better predict
|
||||
- complete logger
|
||||
- SDE: learn the feature extractor?
|
||||
- Refactor: buffer with numpy array instead of pytorch
|
||||
- Refactor: remove duplicated code for evaluation
|
||||
- plotting? -> zoo
|
||||
|
||||
Later:
|
||||
- get_parameters / set_parameters
|
||||
- CNN policies + normalization
|
||||
- tensorboard support
|
||||
- DQN
|
||||
- TRPO
|
||||
- ACER
|
||||
- DDPG
|
||||
- HER -> use stable-baselines because does not depends on tf?
|
||||
- cf github Roadmap
|
||||
|
|
|
|||
8
setup.py
8
setup.py
|
|
@ -23,6 +23,12 @@ setup(name='torchy_baselines',
|
|||
'sphinx',
|
||||
'sphinx-autobuild',
|
||||
'sphinx-rtd-theme'
|
||||
],
|
||||
'extra': [
|
||||
# For render
|
||||
'opencv-python',
|
||||
# For reading logs
|
||||
'pandas'
|
||||
]
|
||||
},
|
||||
description='Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.',
|
||||
|
|
@ -34,7 +40,7 @@ setup(name='torchy_baselines',
|
|||
license="MIT",
|
||||
long_description="",
|
||||
long_description_content_type='text/markdown',
|
||||
version="0.0.6a",
|
||||
version="0.0.8a0",
|
||||
)
|
||||
|
||||
# python setup.py sdist
|
||||
|
|
|
|||
|
|
@ -1,6 +1,3 @@
|
|||
import os
|
||||
|
||||
import gym
|
||||
import pytest
|
||||
|
||||
from torchy_baselines import PPO
|
||||
|
|
@ -15,4 +12,4 @@ from torchy_baselines import PPO
|
|||
[12, dict(pi=[8])],
|
||||
])
|
||||
def test_flexible_mlp(net_arch):
|
||||
model = PPO('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)
|
||||
_ = PPO('MlpPolicy', 'CartPole-v1', policy_kwargs=dict(net_arch=net_arch), n_steps=100).learn(1000)
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import numpy as np
|
||||
import pytest
|
||||
import torch as th
|
||||
|
||||
from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution,\
|
||||
CategoricalDistribution, TanhBijector
|
||||
from torchy_baselines.common.distributions import DiagGaussianDistribution, TanhBijector, \
|
||||
StateDependentNoiseDistribution
|
||||
from torchy_baselines.common.utils import set_random_seed
|
||||
|
||||
|
||||
# TODO: more tests for the other distributions
|
||||
def test_bijector():
|
||||
|
|
@ -18,3 +20,51 @@ def test_bijector():
|
|||
assert th.max(th.abs(squashed_actions)) <= 1.0
|
||||
# Check the inverse method
|
||||
assert th.isclose(TanhBijector.inverse(squashed_actions), actions).all()
|
||||
|
||||
|
||||
def test_sde_distribution():
|
||||
n_samples = int(5e6)
|
||||
n_features = 2
|
||||
n_actions = 1
|
||||
deterministic_actions = th.ones(n_samples, n_actions) * 0.1
|
||||
state = th.ones(n_samples, n_features) * 0.3
|
||||
dist = StateDependentNoiseDistribution(n_actions, full_std=True, squash_output=False)
|
||||
|
||||
set_random_seed(1)
|
||||
_, log_std = dist.proba_distribution_net(n_features)
|
||||
dist.sample_weights(log_std, batch_size=n_samples)
|
||||
|
||||
actions, _ = dist.proba_distribution(deterministic_actions, log_std, state)
|
||||
|
||||
assert th.allclose(actions.mean(), dist.distribution.mean.mean(), rtol=1e-3)
|
||||
assert th.allclose(actions.std(), dist.distribution.scale.mean(), rtol=1e-3)
|
||||
|
||||
|
||||
N_ACTIONS = 1
|
||||
|
||||
|
||||
# TODO: fix for num action > 1
|
||||
# TODO: analytical form for squashed Gaussian?
|
||||
@pytest.mark.parametrize("dist", [
|
||||
DiagGaussianDistribution(N_ACTIONS),
|
||||
StateDependentNoiseDistribution(N_ACTIONS, squash_output=False),
|
||||
])
|
||||
def test_entropy(dist):
|
||||
# The entropy can be approximated by averaging the negative log likelihood
|
||||
# mean negative log likelihood == differential entropy
|
||||
n_samples = int(5e6)
|
||||
n_features = 3
|
||||
set_random_seed(1)
|
||||
state = th.rand(n_samples, n_features)
|
||||
deterministic_actions = th.rand(n_samples, N_ACTIONS)
|
||||
_, log_std = dist.proba_distribution_net(n_features, log_std_init=th.log(th.tensor(0.2)))
|
||||
|
||||
if isinstance(dist, DiagGaussianDistribution):
|
||||
actions, dist = dist.proba_distribution(deterministic_actions, log_std)
|
||||
else:
|
||||
dist.sample_weights(log_std, batch_size=n_samples)
|
||||
actions, dist = dist.proba_distribution(deterministic_actions, log_std, state)
|
||||
|
||||
entropy = dist.entropy()
|
||||
log_prob = dist.log_prob(actions)
|
||||
assert th.allclose(entropy.mean(), -log_prob.mean(), rtol=5e-3)
|
||||
|
|
|
|||
82
tests/test_logger.py
Normal file
82
tests/test_logger.py
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
import os
|
||||
import shutil
|
||||
|
||||
import pytest
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.logger import (make_output_format, read_csv, read_json, DEBUG, ScopedConfigure,
|
||||
info, debug, set_level, configure, logkv, logkvs, dumpkvs, logkv_mean, warn, error, reset)
|
||||
|
||||
KEY_VALUES = {
|
||||
"test": 1,
|
||||
"b": -3.14,
|
||||
"8": 9.9,
|
||||
"l": [1, 2],
|
||||
"a": np.array([1, 2, 3]),
|
||||
"f": np.array(1),
|
||||
"g": np.array([[[1]]]),
|
||||
}
|
||||
|
||||
LOG_DIR = '/tmp/torchy_baselines/'
|
||||
|
||||
|
||||
def test_main():
|
||||
"""
|
||||
tests for the logger module
|
||||
"""
|
||||
info("hi")
|
||||
debug("shouldn't appear")
|
||||
set_level(DEBUG)
|
||||
debug("should appear")
|
||||
folder = "/tmp/testlogging"
|
||||
if os.path.exists(folder):
|
||||
shutil.rmtree(folder)
|
||||
configure(folder=folder)
|
||||
logkv("a", 3)
|
||||
logkv("b", 2.5)
|
||||
dumpkvs()
|
||||
logkv("b", -2.5)
|
||||
logkv("a", 5.5)
|
||||
dumpkvs()
|
||||
info("^^^ should see a = 5.5")
|
||||
logkv_mean("b", -22.5)
|
||||
logkv_mean("b", -44.4)
|
||||
logkv("a", 5.5)
|
||||
dumpkvs()
|
||||
with ScopedConfigure(None, None):
|
||||
info("^^^ should see b = 33.3")
|
||||
|
||||
with ScopedConfigure("/tmp/test-logger/", ["json"]):
|
||||
logkv("b", -2.5)
|
||||
dumpkvs()
|
||||
|
||||
reset()
|
||||
logkv("a", "longasslongasslongasslongasslongasslongassvalue")
|
||||
dumpkvs()
|
||||
warn("hey")
|
||||
error("oh")
|
||||
logkvs({"test": 1})
|
||||
|
||||
|
||||
@pytest.mark.parametrize('_format', ['stdout', 'log', 'json', 'csv'])
|
||||
def test_make_output(_format):
|
||||
"""
|
||||
test make output
|
||||
|
||||
:param _format: (str) output format
|
||||
"""
|
||||
writer = make_output_format(_format, LOG_DIR)
|
||||
writer.writekvs(KEY_VALUES)
|
||||
if _format == "csv":
|
||||
read_csv(LOG_DIR + 'progress.csv')
|
||||
elif _format == 'json':
|
||||
read_json(LOG_DIR + 'progress.json')
|
||||
writer.close()
|
||||
|
||||
|
||||
def test_make_output_fail():
|
||||
"""
|
||||
test value error on logger
|
||||
"""
|
||||
with pytest.raises(ValueError):
|
||||
make_output_format('dummy_format', LOG_DIR)
|
||||
|
|
@ -37,9 +37,10 @@ def test_onpolicy(model_class, env_id):
|
|||
os.remove("test_save.zip")
|
||||
|
||||
|
||||
def test_sac():
|
||||
@pytest.mark.parametrize("ent_coef", ['auto', 0.01])
|
||||
def test_sac(ent_coef):
|
||||
model = SAC('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]),
|
||||
learning_starts=100, verbose=1, create_eval_env=True, ent_coef='auto',
|
||||
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")
|
||||
|
|
|
|||
|
|
@ -1,13 +1,12 @@
|
|||
import numpy as np
|
||||
import os
|
||||
import pytest
|
||||
from copy import deepcopy
|
||||
import numpy as np
|
||||
|
||||
import torch as th
|
||||
from copy import deepcopy
|
||||
|
||||
from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
|
||||
from torchy_baselines.common.identity_env import IdentityEnvBox
|
||||
from torchy_baselines.common.vec_env import DummyVecEnv
|
||||
from torchy_baselines.common.identity_env import IdentityEnvBox, IdentityEnv
|
||||
|
||||
MODEL_LIST = [
|
||||
CEMRL,
|
||||
|
|
@ -81,7 +80,8 @@ def test_save_load(model_class):
|
|||
for optimizer, opt_state in opt_params.items():
|
||||
for param_group_idx, param_group in enumerate(opt_state['param_groups']):
|
||||
for param_key, param_value in param_group.items():
|
||||
if param_key == 'params': # don't know how to handle params correctly, therefore only check if we have the same amount
|
||||
# don't know how to handle params correctly, therefore only check if we have the same amount
|
||||
if param_key == 'params':
|
||||
assert len(param_value) == len(
|
||||
new_opt_params[optimizer]['param_groups'][param_group_idx][param_key])
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -1,69 +1,65 @@
|
|||
import pytest
|
||||
|
||||
import gym
|
||||
import torch as th
|
||||
from torch.distributions import Normal
|
||||
|
||||
from torchy_baselines import A2C, TD3
|
||||
from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize
|
||||
from torchy_baselines.common.monitor import Monitor
|
||||
from torchy_baselines import A2C, TD3, SAC
|
||||
|
||||
|
||||
def test_state_dependent_exploration():
|
||||
def test_state_dependent_exploration_grad():
|
||||
"""
|
||||
Check that the gradient correspond to the expected one
|
||||
"""
|
||||
n_states = 2
|
||||
state_dim = 3
|
||||
# TODO: fix for action_dim > 1
|
||||
action_dim = 1
|
||||
sigma = th.ones(state_dim, 1, requires_grad=True)
|
||||
action_dim = 10
|
||||
sigma_hat = th.ones(state_dim, action_dim, requires_grad=True)
|
||||
# Reduce the number of parameters
|
||||
# sigma_ = th.ones(state_dim, action_dim) * sigma_
|
||||
|
||||
# weights_dist = Normal(th.zeros_like(log_sigma), th.exp(log_sigma))
|
||||
th.manual_seed(2)
|
||||
weights_dist = Normal(th.zeros_like(sigma), sigma)
|
||||
|
||||
weights_dist = Normal(th.zeros_like(sigma_hat), sigma_hat)
|
||||
weights = weights_dist.rsample()
|
||||
|
||||
state = th.rand(n_states, state_dim)
|
||||
mu = th.ones(action_dim)
|
||||
# print(weights.shape, state.shape)
|
||||
noise = th.mm(state, weights)
|
||||
|
||||
variance = th.mm(state ** 2, sigma ** 2)
|
||||
action = mu + noise
|
||||
|
||||
variance = th.mm(state ** 2, sigma_hat ** 2)
|
||||
action_dist = Normal(mu, th.sqrt(variance))
|
||||
|
||||
loss = action_dist.log_prob((mu + noise).detach()).mean()
|
||||
# Sum over the action dimension because we assume they are independent
|
||||
loss = action_dist.log_prob(action.detach()).sum(dim=-1).mean()
|
||||
loss.backward()
|
||||
|
||||
# From Rueckstiess paper
|
||||
grad = th.zeros_like(sigma)
|
||||
# From Rueckstiess paper: check that the computed gradient
|
||||
# correspond to the analytical form
|
||||
grad = th.zeros_like(sigma_hat)
|
||||
for j in range(action_dim):
|
||||
# sigma_hat is the std of the gaussian distribution of the noise matrix weights
|
||||
# sigma_j = sum_j(state_i **2 * sigma_hat_ij ** 2)
|
||||
# sigma_j is the standard deviation of the policy gaussian distribution
|
||||
sigma_j = th.sqrt(variance[:, j])
|
||||
for i in range(state_dim):
|
||||
a = ((noise[:, j] ** 2 - variance[:, j]) / (variance[:, j] ** 2)) * (state[:, i] ** 2 * sigma[i, j])
|
||||
grad[i, j] = a.mean()
|
||||
# Derivative of the log probability of the jth component of the action
|
||||
# w.r.t. the standard deviation sigma_j
|
||||
d_log_policy_j = (noise[:, j] ** 2 - sigma_j ** 2) / sigma_j ** 3
|
||||
# Derivative of sigma_j w.r.t. sigma_hat_ij
|
||||
d_log_sigma_j = (state[:, i] ** 2 * sigma_hat[i, j]) / sigma_j
|
||||
# Chain rule, average over the minibatch
|
||||
grad[i, j] = (d_log_policy_j * d_log_sigma_j).mean()
|
||||
|
||||
# sigma.grad should be equal to grad
|
||||
assert sigma.grad.allclose(grad)
|
||||
assert sigma_hat.grad.allclose(grad)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_class", [A2C])
|
||||
def test_state_dependent_noise(model_class):
|
||||
env_id = 'MountainCarContinuous-v0'
|
||||
|
||||
env = VecNormalize(DummyVecEnv([lambda: Monitor(gym.make(env_id))]), norm_reward=True)
|
||||
eval_env = VecNormalize(DummyVecEnv([lambda: Monitor(gym.make(env_id))]), training=False, norm_reward=False)
|
||||
|
||||
model = model_class('MlpPolicy', env, n_steps=200, use_sde=True, ent_coef=0.00, verbose=1, learning_rate=3e-4,
|
||||
policy_kwargs=dict(log_std_init=0.0, ortho_init=False), seed=None)
|
||||
model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_class", [TD3])
|
||||
def test_state_dependent_offpolicy_noise(model_class):
|
||||
@pytest.mark.parametrize("model_class", [TD3, SAC, A2C])
|
||||
@pytest.mark.parametrize("sde_net_arch", [None, [32, 16], []])
|
||||
def test_state_dependent_offpolicy_noise(model_class, sde_net_arch):
|
||||
model = model_class('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True,
|
||||
verbose=1, policy_kwargs=dict(log_std_init=-2))
|
||||
verbose=1, policy_kwargs=dict(log_std_init=-2, sde_net_arch=sde_net_arch))
|
||||
model.learn(total_timesteps=int(1000), eval_freq=500)
|
||||
|
||||
|
||||
|
|
@ -72,6 +68,6 @@ def test_scheduler():
|
|||
return -2.0 * progress + 1
|
||||
|
||||
model = TD3('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True,
|
||||
verbose=1, sde_log_std_scheduler=scheduler)
|
||||
verbose=1, sde_log_std_scheduler=scheduler)
|
||||
model.learn(total_timesteps=int(1000), eval_freq=500)
|
||||
assert th.isclose(model.actor.log_std, th.ones_like(model.actor.log_std)).all()
|
||||
|
|
|
|||
|
|
@ -325,9 +325,8 @@ def test_vecenv_wrapper_getattr():
|
|||
assert wrapped.name_test() == CustomWrapperBB
|
||||
|
||||
double_wrapped = CustomWrapperA(CustomWrapperB(wrapped))
|
||||
dummy = double_wrapped.var_a # should not raise as it is directly defined here
|
||||
_ = double_wrapped.var_a # should not raise as it is directly defined here
|
||||
with pytest.raises(AttributeError): # should raise due to ambiguity
|
||||
dummy = double_wrapped.var_b
|
||||
_ = double_wrapped.var_b
|
||||
with pytest.raises(AttributeError): # should raise as does not exist
|
||||
dummy = double_wrapped.nonexistent_attribute
|
||||
del dummy # keep linter happy
|
||||
_ = double_wrapped.nonexistent_attribute
|
||||
|
|
|
|||
|
|
@ -3,12 +3,49 @@ import pytest
|
|||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.running_mean_std import RunningMeanStd
|
||||
from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv
|
||||
from torchy_baselines.common.vec_env.vec_normalize import VecNormalize
|
||||
from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize, VecFrameStack, sync_envs_normalization
|
||||
from torchy_baselines import CEMRL, SAC, TD3
|
||||
|
||||
ENV_ID = 'Pendulum-v0'
|
||||
|
||||
def make_env():
|
||||
return gym.make(ENV_ID)
|
||||
|
||||
def check_rms_equal(rmsa, rmsb):
|
||||
assert np.all(rmsa.mean == rmsb.mean)
|
||||
assert np.all(rmsa.var == rmsb.var)
|
||||
assert np.all(rmsa.count == rmsb.count)
|
||||
|
||||
|
||||
def check_vec_norm_equal(norma, normb):
|
||||
assert norma.observation_space == normb.observation_space
|
||||
assert norma.action_space == normb.action_space
|
||||
assert norma.num_envs == normb.num_envs
|
||||
|
||||
check_rms_equal(norma.obs_rms, normb.obs_rms)
|
||||
check_rms_equal(norma.ret_rms, normb.ret_rms)
|
||||
assert norma.clip_obs == normb.clip_obs
|
||||
assert norma.clip_reward == normb.clip_reward
|
||||
assert norma.norm_obs == normb.norm_obs
|
||||
assert norma.norm_reward == normb.norm_reward
|
||||
|
||||
assert np.all(norma.ret == normb.ret)
|
||||
assert norma.gamma == normb.gamma
|
||||
assert norma.epsilon == normb.epsilon
|
||||
assert norma.training == normb.training
|
||||
|
||||
def _make_warmstart_cartpole():
|
||||
"""Warm-start VecNormalize by stepping through CartPole"""
|
||||
venv = DummyVecEnv([lambda: gym.make("CartPole-v1")])
|
||||
venv = VecNormalize(venv)
|
||||
venv.reset()
|
||||
venv.get_original_obs()
|
||||
|
||||
for _ in range(100):
|
||||
actions = [venv.action_space.sample()]
|
||||
venv.step(actions)
|
||||
return venv
|
||||
|
||||
|
||||
def test_runningmeanstd():
|
||||
"""Test RunningMeanStd object"""
|
||||
|
|
@ -27,29 +64,87 @@ def test_runningmeanstd():
|
|||
assert np.allclose(moments_1, moments_2)
|
||||
|
||||
|
||||
def test_vec_env():
|
||||
def test_vec_env(tmpdir):
|
||||
"""Test VecNormalize Object"""
|
||||
clip_obs = 0.5
|
||||
clip_reward = 5.0
|
||||
|
||||
def make_env():
|
||||
return gym.make(ENV_ID)
|
||||
|
||||
env = DummyVecEnv([make_env])
|
||||
env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.)
|
||||
_, done = env.reset(), [False]
|
||||
obs = None
|
||||
orig_venv = DummyVecEnv([make_env])
|
||||
norm_venv = VecNormalize(orig_venv, norm_obs=True, norm_reward=True, clip_obs=clip_obs, clip_reward=clip_reward)
|
||||
_, done = norm_venv.reset(), [False]
|
||||
while not done[0]:
|
||||
actions = [env.action_space.sample()]
|
||||
obs, _, done, _ = env.step(actions)
|
||||
assert np.max(obs) <= 10
|
||||
actions = [norm_venv.action_space.sample()]
|
||||
obs, rew, done, _ = norm_venv.step(actions)
|
||||
assert np.max(np.abs(obs)) <= clip_obs
|
||||
assert np.max(np.abs(rew)) <= clip_reward
|
||||
|
||||
path = str(tmpdir.join("vec_normalize"))
|
||||
norm_venv.save(path)
|
||||
deserialized = VecNormalize.load(path, venv=orig_venv)
|
||||
check_vec_norm_equal(norm_venv, deserialized)
|
||||
|
||||
|
||||
def test_get_original():
|
||||
venv = _make_warmstart_cartpole()
|
||||
for _ in range(3):
|
||||
actions = [venv.action_space.sample()]
|
||||
obs, rewards, _, _ = venv.step(actions)
|
||||
obs = obs[0]
|
||||
orig_obs = venv.get_original_obs()[0]
|
||||
rewards = rewards[0]
|
||||
orig_rewards = venv.get_original_reward()[0]
|
||||
|
||||
assert np.all(orig_rewards == 1)
|
||||
assert orig_obs.shape == obs.shape
|
||||
assert orig_rewards.dtype == rewards.dtype
|
||||
assert not np.array_equal(orig_obs, obs)
|
||||
assert not np.array_equal(orig_rewards, rewards)
|
||||
np.testing.assert_allclose(venv.normalize_obs(orig_obs), obs)
|
||||
np.testing.assert_allclose(venv.normalize_reward(orig_rewards), rewards)
|
||||
|
||||
|
||||
def test_normalize_external():
|
||||
venv = _make_warmstart_cartpole()
|
||||
|
||||
rewards = np.array([1, 1])
|
||||
norm_rewards = venv.normalize_reward(rewards)
|
||||
assert norm_rewards.shape == rewards.shape
|
||||
# Episode return is almost always >= 1 in CartPole. So reward should shrink.
|
||||
assert np.all(norm_rewards < 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_class", [SAC, TD3, CEMRL])
|
||||
def test_offpolicy_normalization(model_class):
|
||||
env = DummyVecEnv([lambda: gym.make(ENV_ID)])
|
||||
env = DummyVecEnv([make_env])
|
||||
env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.)
|
||||
|
||||
eval_env = DummyVecEnv([lambda: gym.make(ENV_ID)])
|
||||
eval_env = VecNormalize(eval_env, norm_obs=True, norm_reward=False, clip_obs=10., clip_reward=10.)
|
||||
eval_env = DummyVecEnv([make_env])
|
||||
eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=False, clip_obs=10., clip_reward=10.)
|
||||
|
||||
model = model_class('MlpPolicy', env, verbose=1)
|
||||
model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500)
|
||||
|
||||
|
||||
def test_sync_vec_normalize():
|
||||
env = DummyVecEnv([make_env])
|
||||
env = VecNormalize(env, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.)
|
||||
env = VecFrameStack(env, 1)
|
||||
|
||||
eval_env = DummyVecEnv([make_env])
|
||||
eval_env = VecNormalize(eval_env, training=False, norm_obs=True, norm_reward=True, clip_obs=10., clip_reward=10.)
|
||||
eval_env = VecFrameStack(eval_env, 1)
|
||||
|
||||
env.reset()
|
||||
# Initialize running mean
|
||||
for _ in range(100):
|
||||
env.step([env.action_space.sample()])
|
||||
|
||||
obs = env.reset()
|
||||
original_obs = env.get_original_obs()
|
||||
# Normalization must be different
|
||||
assert not np.allclose(obs, eval_env.normalize_obs(original_obs))
|
||||
|
||||
sync_envs_normalization(env, eval_env)
|
||||
|
||||
# Now they must be synced
|
||||
assert np.allclose(obs, eval_env.normalize_obs(original_obs))
|
||||
|
|
|
|||
|
|
@ -4,4 +4,4 @@ from torchy_baselines.ppo import PPO
|
|||
from torchy_baselines.sac import SAC
|
||||
from torchy_baselines.td3 import TD3
|
||||
|
||||
__version__ = "0.0.6a"
|
||||
__version__ = "0.0.8a0"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ import torch.nn.functional as F
|
|||
|
||||
from torchy_baselines.common.utils import explained_variance
|
||||
from torchy_baselines.ppo.ppo import PPO
|
||||
from torchy_baselines.ppo.policies import PPOPolicy
|
||||
from torchy_baselines.common import logger
|
||||
|
||||
|
||||
|
|
@ -34,6 +33,8 @@ class A2C(PPO):
|
|||
:param use_rms_prop: (bool) Whether to use RMSprop (default) or Adam as optimizer
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
|
||||
instead of action noise exploration (default: False)
|
||||
:param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
|
||||
Default: -1 (only sample at the beginning of the rollout)
|
||||
:param normalize_advantage: (bool) Whether to normalize or not the advantage
|
||||
:param tensorboard_log: (str) the log location for tensorboard (if None, no logging)
|
||||
:param create_eval_env: (bool) Whether to create a second environment that will be
|
||||
|
|
@ -45,11 +46,10 @@ class A2C(PPO):
|
|||
Setting it to auto, the code will be run on the GPU if possible.
|
||||
:param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance
|
||||
"""
|
||||
|
||||
def __init__(self, policy, env, learning_rate=7e-4,
|
||||
n_steps=5, gamma=0.99, gae_lambda=1.0,
|
||||
ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5,
|
||||
rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False,
|
||||
rms_prop_eps=1e-5, use_rms_prop=True, use_sde=False, sde_sample_freq=-1,
|
||||
normalize_advantage=False, tensorboard_log=None, create_eval_env=False,
|
||||
policy_kwargs=None, verbose=0, seed=0, device='auto',
|
||||
_init_setup_model=True):
|
||||
|
|
@ -57,7 +57,8 @@ class A2C(PPO):
|
|||
super(A2C, self).__init__(policy, env, learning_rate=learning_rate,
|
||||
n_steps=n_steps, batch_size=None, n_epochs=1,
|
||||
gamma=gamma, gae_lambda=gae_lambda, ent_coef=ent_coef,
|
||||
vf_coef=vf_coef, max_grad_norm=max_grad_norm, use_sde=use_sde,
|
||||
vf_coef=vf_coef, max_grad_norm=max_grad_norm,
|
||||
use_sde=use_sde, sde_sample_freq=sde_sample_freq,
|
||||
tensorboard_log=tensorboard_log, policy_kwargs=policy_kwargs,
|
||||
verbose=verbose, device=device, create_eval_env=create_eval_env,
|
||||
seed=seed, _init_setup_model=False)
|
||||
|
|
@ -113,6 +114,7 @@ class A2C(PPO):
|
|||
# Optimization step
|
||||
self.policy.optimizer.zero_grad()
|
||||
loss.backward()
|
||||
|
||||
# Clip grad norm
|
||||
th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
|
||||
self.policy.optimizer.step()
|
||||
|
|
|
|||
|
|
@ -155,11 +155,12 @@ class CEMRL(TD3):
|
|||
self.actor.load_from_vector(self.es.mu)
|
||||
sync_envs_normalization(self.env, eval_env)
|
||||
|
||||
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
|
||||
mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes)
|
||||
evaluations.append(mean_reward)
|
||||
|
||||
if self.verbose > 0:
|
||||
print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
|
||||
print("Eval num_timesteps={}, "
|
||||
"episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward))
|
||||
print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
|
||||
|
||||
actor_steps = 0
|
||||
|
|
|
|||
|
|
@ -9,11 +9,12 @@ import gym
|
|||
import torch as th
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common import logger
|
||||
from torchy_baselines.common.policies import get_policy_from_name
|
||||
from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
|
||||
from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize
|
||||
from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, sync_envs_normalization
|
||||
from torchy_baselines.common.monitor import Monitor
|
||||
from torchy_baselines.common import logger
|
||||
from torchy_baselines.common.evaluation import evaluate_policy
|
||||
from torchy_baselines.common.save_util import data_to_json, json_to_data
|
||||
|
||||
|
||||
|
|
@ -37,12 +38,17 @@ class BaseRLModel(object):
|
|||
:param monitor_wrapper: (bool) When creating an environment, whether to wrap it
|
||||
or not in a Monitor wrapper.
|
||||
:param seed: (int) Seed for the pseudo random generators
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
|
||||
instead of action noise exploration (default: False)
|
||||
:param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
|
||||
Default: -1 (only sample at the beginning of the rollout)
|
||||
"""
|
||||
__metaclass__ = ABCMeta
|
||||
|
||||
def __init__(self, policy, env, policy_base, policy_kwargs=None,
|
||||
verbose=0, device='auto', support_multi_env=False,
|
||||
create_eval_env=False, monitor_wrapper=True, seed=None):
|
||||
create_eval_env=False, monitor_wrapper=True, seed=None,
|
||||
use_sde=False, sde_sample_freq=-1):
|
||||
if isinstance(policy, str) and policy_base is not None:
|
||||
self.policy_class = get_policy_from_name(policy_base, policy)
|
||||
else:
|
||||
|
|
@ -70,7 +76,9 @@ class BaseRLModel(object):
|
|||
self.action_noise = None
|
||||
# Used for SDE only
|
||||
self.rollout_data = None
|
||||
self.use_sde = False
|
||||
self.on_policy_exploration = False
|
||||
self.use_sde = use_sde
|
||||
self.sde_sample_freq = sde_sample_freq
|
||||
# Track the training progress (from 1 to 0)
|
||||
# this is used to update the learning rate
|
||||
self._current_progress = 1
|
||||
|
|
@ -217,7 +225,8 @@ class BaseRLModel(object):
|
|||
:param env: (gym.Env) The environment for learning a policy
|
||||
"""
|
||||
if self.check_env(env, self.observation_space, self.action_space) is False:
|
||||
raise ValueError("The given environment is not compatible with model: observation and action spaces do not match")
|
||||
raise ValueError("The given environment is not compatible with model: "
|
||||
"observation and action spaces do not match")
|
||||
# it must be coherent now
|
||||
# if it is not a VecEnv, make it a VecEnv
|
||||
if not isinstance(env, VecEnv):
|
||||
|
|
@ -250,24 +259,6 @@ class BaseRLModel(object):
|
|||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def pretrain(self, dataset, n_epochs=10, learning_rate=1e-4,
|
||||
adam_epsilon=1e-8, val_interval=None):
|
||||
"""
|
||||
Pretrain a model using behavior cloning:
|
||||
supervised learning given an expert dataset.
|
||||
|
||||
NOTE: only Box and Discrete spaces are supported for now.
|
||||
|
||||
:param dataset: (ExpertDataset) Dataset manager
|
||||
:param n_epochs: (int) Number of iterations on the training set
|
||||
:param learning_rate: (float) Learning rate
|
||||
:param adam_epsilon: (float) the epsilon value for the adam optimizer
|
||||
:param val_interval: (int) Report training and validation losses every n epochs.
|
||||
By default, every 10th of the maximum number of epochs.
|
||||
:return: (BaseRLModel) the pretrained model
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="run",
|
||||
eval_env=None, eval_freq=-1, n_eval_episodes=5, reset_num_timesteps=True):
|
||||
|
|
@ -280,12 +271,12 @@ class BaseRLModel(object):
|
|||
:param log_interval: (int) The number of timesteps before logging.
|
||||
:param tb_log_name: (str) the name of the run for tensorboard log
|
||||
:param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging)
|
||||
:param eval_env: (gym.Env)
|
||||
:param eval_freq: (int)
|
||||
:param n_eval_episodes: (int)
|
||||
:param eval_env: (gym.Env) Environment that will be used to evaluate the agent
|
||||
:param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little)
|
||||
:param n_eval_episodes: (int) Number of episode to evaluate the agent
|
||||
:return: (BaseRLModel) the trained model
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def predict(self, observation, state=None, mask=None, deterministic=False):
|
||||
|
|
@ -298,13 +289,14 @@ class BaseRLModel(object):
|
|||
:param deterministic: (bool) Whether or not to return deterministic actions.
|
||||
:return: (np.ndarray, np.ndarray) the model's action and the next state (used in recurrent policies)
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
def load_parameters(self, load_dict, opt_params=None):
|
||||
def load_parameters(self, load_dict, opt_params):
|
||||
"""
|
||||
Load model parameters from a dictionary
|
||||
load_dict should contain all keys from torch.model.state_dict()
|
||||
If opt_params are given this does also load agent's optimizer-parameters, but can only be handled in child classes.
|
||||
If opt_params are given this does also load agent's optimizer-parameters,
|
||||
but can only be handled in child classes.
|
||||
|
||||
|
||||
:param load_dict: (dict) dict of parameters from model.state_dict()
|
||||
|
|
@ -342,6 +334,7 @@ class BaseRLModel(object):
|
|||
env = data["env"]
|
||||
|
||||
# first create model, but only setup if a env was given
|
||||
# noinspection PyArgumentList
|
||||
model = cls(policy=data["policy_class"], env=env, _init_setup_model=env is not None)
|
||||
|
||||
# load parameters
|
||||
|
|
@ -357,7 +350,8 @@ class BaseRLModel(object):
|
|||
:param load_path: (str) Where to load the model from
|
||||
:param load_data: (bool) Whether we should load and return data
|
||||
(class parameters). Mainly used by 'load_parameters' to only load model parameters (weights)
|
||||
:return: (dict),(dict),(dict) Class parameters, model parameters (state_dict) and dict of optimizer parameters (dict of state_dict)
|
||||
:return: (dict),(dict),(dict) Class parameters, model parameters (state_dict)
|
||||
and dict of optimizer parameters (dict of state_dict)
|
||||
"""
|
||||
# Check if file exists if load_path is a string
|
||||
if isinstance(load_path, str):
|
||||
|
|
@ -395,7 +389,7 @@ class BaseRLModel(object):
|
|||
|
||||
# check for all other .pth files
|
||||
other_files = [file_name for file_name in namelist if
|
||||
os.path.splitext(file_name)[1] == ".pth" and file_name != "params.pth"]
|
||||
os.path.splitext(file_name)[1] == ".pth" and file_name != "params.pth"]
|
||||
# if there are any other files which end with .pth and aren't "params.pth"
|
||||
# assume that they each are optimizer parameters
|
||||
if len(other_files) > 0:
|
||||
|
|
@ -503,7 +497,8 @@ class BaseRLModel(object):
|
|||
if self.use_sde:
|
||||
self.actor.reset_noise()
|
||||
# Reset rollout data
|
||||
self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']}
|
||||
if self.on_policy_exploration:
|
||||
self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones', 'values']}
|
||||
|
||||
while total_steps < n_steps or total_episodes < n_episodes:
|
||||
done = False
|
||||
|
|
@ -512,7 +507,13 @@ class BaseRLModel(object):
|
|||
episode_reward, episode_timesteps = 0.0, 0
|
||||
|
||||
while not done:
|
||||
if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
|
||||
# Sample a new noise matrix
|
||||
self.actor.reset_noise()
|
||||
|
||||
# Select action randomly or according to policy
|
||||
# TODO: use action from policy when using SDE during the warmup phase?
|
||||
# if num_timesteps < learning_starts and not self.use_sde:
|
||||
if num_timesteps < learning_starts:
|
||||
# Warmup phase
|
||||
unscaled_action = np.array([self.action_space.sample()])
|
||||
|
|
@ -562,6 +563,8 @@ class BaseRLModel(object):
|
|||
self.rollout_data['actions'].append(scaled_action[0].copy())
|
||||
self.rollout_data['rewards'].append(reward[0].copy())
|
||||
self.rollout_data['dones'].append(np.array(done_bool[0]).copy())
|
||||
obs_tensor = th.FloatTensor(obs).to(self.device)
|
||||
self.rollout_data['values'].append(self.vf_net(obs_tensor)[0].cpu().detach().numpy())
|
||||
|
||||
obs = new_obs
|
||||
# Save the true unnormalized observation
|
||||
|
|
@ -580,6 +583,7 @@ class BaseRLModel(object):
|
|||
total_episodes += 1
|
||||
episode_rewards.append(episode_reward)
|
||||
total_timesteps.append(episode_timesteps)
|
||||
# TODO: reset SDE matrix at the end of the episode?
|
||||
if action_noise is not None:
|
||||
action_noise.reset()
|
||||
|
||||
|
|
@ -603,19 +607,24 @@ class BaseRLModel(object):
|
|||
|
||||
# Post processing
|
||||
if self.rollout_data is not None:
|
||||
for key in ['observations', 'actions', 'rewards', 'dones']:
|
||||
for key in ['observations', 'actions', 'rewards', 'dones', 'values']:
|
||||
self.rollout_data[key] = th.FloatTensor(np.array(self.rollout_data[key])).to(self.device)
|
||||
|
||||
self.rollout_data['returns'] = self.rollout_data['rewards'].clone()
|
||||
# Compute return
|
||||
self.rollout_data['advantage'] = self.rollout_data['rewards'].clone()
|
||||
|
||||
# Compute return and advantage
|
||||
last_return = 0.0
|
||||
for step in reversed(range(len(self.rollout_data['rewards']))):
|
||||
if step == len(self.rollout_data['rewards']) - 1:
|
||||
last_return = self.rollout_data['rewards'][step]
|
||||
next_non_terminal = 1.0 - done[0]
|
||||
next_value = self.vf_net(th.FloatTensor(obs).to(self.device))[0].detach()
|
||||
last_return = self.rollout_data['rewards'][step] + next_non_terminal * next_value
|
||||
else:
|
||||
next_non_terminal = 1.0 - self.rollout_data['dones'][step + 1]
|
||||
last_return = self.rollout_data['rewards'][step] + self.gamma * last_return * next_non_terminal
|
||||
self.rollout_data['returns'][step] = last_return
|
||||
self.rollout_data['advantage'] = self.rollout_data['returns'] - self.rollout_data['values']
|
||||
|
||||
return mean_reward, total_steps, total_episodes, obs
|
||||
|
||||
|
|
@ -652,9 +661,9 @@ class BaseRLModel(object):
|
|||
with archive.open('params.pth', mode="w") as param_file:
|
||||
th.save(params, param_file)
|
||||
if opt_params is not None:
|
||||
for file_name, dict in opt_params.items():
|
||||
for file_name, dict_ in opt_params.items():
|
||||
with archive.open(file_name + '.pth', mode="w") as opt_param_file:
|
||||
th.save(dict, opt_param_file)
|
||||
th.save(dict_, opt_param_file)
|
||||
|
||||
@staticmethod
|
||||
def excluded_save_params():
|
||||
|
|
@ -664,7 +673,7 @@ class BaseRLModel(object):
|
|||
|
||||
:return: ([str]) List of parameters that should be excluded from save
|
||||
"""
|
||||
return ["env", "eval_env", "replay_buffer", "rollout_buffer"]
|
||||
return ["env", "eval_env", "replay_buffer", "rollout_buffer", "_vec_normalize_env"]
|
||||
|
||||
def save(self, path, exclude=None, include=None):
|
||||
"""
|
||||
|
|
@ -694,3 +703,26 @@ class BaseRLModel(object):
|
|||
params_to_save = self.get_policy_parameters()
|
||||
opt_params_to_save = self.get_opt_parameters()
|
||||
self._save_to_file_zip(path, data=data, params=params_to_save, opt_params=opt_params_to_save)
|
||||
|
||||
def _eval_policy(self, eval_freq, eval_env, n_eval_episodes,
|
||||
timesteps_since_eval, deterministic=True):
|
||||
"""
|
||||
Evaluate the current policy on a test environment.
|
||||
|
||||
:param eval_env: (gym.Env) Environment that will be used to evaluate the agent
|
||||
:param eval_freq: (int) Evaluate the agent every `eval_freq` timesteps (this may vary a little)
|
||||
:param n_eval_episodes: (int) Number of episode to evaluate the agent
|
||||
:parma timesteps_since_eval: (int) Number of timesteps since last evaluation
|
||||
:param deterministic: (bool) Whether to use deterministic or stochastic actions
|
||||
:return: (int) Number of timesteps since last evaluation
|
||||
"""
|
||||
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
|
||||
timesteps_since_eval %= eval_freq
|
||||
# Synchronise the normalization stats if needed
|
||||
sync_envs_normalization(self.env, eval_env)
|
||||
mean_reward, std_reward = evaluate_policy(self, eval_env, n_eval_episodes, deterministic=deterministic)
|
||||
if self.verbose > 0:
|
||||
print("Eval num_timesteps={}, "
|
||||
"episode_reward={:.2f} +/- {:.2f}".format(self.num_timesteps, mean_reward, std_reward))
|
||||
print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
|
||||
return timesteps_since_eval
|
||||
|
|
|
|||
|
|
@ -1,8 +1,6 @@
|
|||
import numpy as np
|
||||
import torch as th
|
||||
|
||||
from torchy_baselines.common.vec_env import unwrap_vec_normalize
|
||||
|
||||
|
||||
class BaseBuffer(object):
|
||||
"""
|
||||
|
|
@ -79,7 +77,8 @@ class BaseBuffer(object):
|
|||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def _normalize_obs(self, obs, env=None):
|
||||
@staticmethod
|
||||
def _normalize_obs(obs, env=None):
|
||||
if env is not None:
|
||||
# TODO: get rid of pytorch - numpy conversion
|
||||
return th.FloatTensor(env.normalize_obs(obs.numpy()))
|
||||
|
|
|
|||
|
|
@ -1,8 +1,6 @@
|
|||
import numpy as np
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
from torch.distributions import Normal, Categorical
|
||||
import torch.nn.functional as F
|
||||
from gym import spaces
|
||||
|
||||
|
||||
|
|
@ -14,17 +12,8 @@ class Distribution(object):
|
|||
"""
|
||||
returns the log likelihood
|
||||
|
||||
:param x: (str) the labels of each index
|
||||
:return: ([float]) The log likelihood of the distribution
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def kl_div(self, other):
|
||||
"""
|
||||
Calculates the Kullback-Leibler divergence from the given probabilty distribution
|
||||
|
||||
:param other: ([float]) the distribution to compare with
|
||||
:return: (float) the KL divergence of the two distributions
|
||||
:param x: (object) the taken action
|
||||
:return: (th.Tensor) The log likelihood of the distribution
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
|
@ -32,7 +21,7 @@ class Distribution(object):
|
|||
"""
|
||||
Returns shannon's entropy of the probability
|
||||
|
||||
:return: (float) the entropy
|
||||
:return: (th.Tensor) the entropy
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
|
@ -40,7 +29,7 @@ class Distribution(object):
|
|||
"""
|
||||
returns a sample from the probabilty distribution
|
||||
|
||||
:return: (Tensorflow Tensor) the stochastic action
|
||||
:return: (th.Tensor) the stochastic action
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
|
@ -52,6 +41,7 @@ class DiagGaussianDistribution(Distribution):
|
|||
|
||||
:param action_dim: (int) Number of continuous actions
|
||||
"""
|
||||
|
||||
def __init__(self, action_dim):
|
||||
super(DiagGaussianDistribution, self).__init__()
|
||||
self.distribution = None
|
||||
|
|
@ -138,6 +128,7 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
|
|||
:param action_dim: (int) Number of continuous actions
|
||||
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
|
||||
"""
|
||||
|
||||
def __init__(self, action_dim, epsilon=1e-6):
|
||||
super(SquashedDiagGaussianDistribution, self).__init__(action_dim)
|
||||
# Avoid NaN (prevents division by zero or log of zero)
|
||||
|
|
@ -185,6 +176,7 @@ class CategoricalDistribution(Distribution):
|
|||
|
||||
:param action_dim: (int) Number of discrete actions
|
||||
"""
|
||||
|
||||
def __init__(self, action_dim):
|
||||
super(CategoricalDistribution, self).__init__()
|
||||
self.distribution = None
|
||||
|
|
@ -243,23 +235,28 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
|
||||
:param squash_output: (bool) Whether to squash the output using a tanh function,
|
||||
this allows to ensure boundaries.
|
||||
:param learn_features: (bool) Whether to learn features for SDE or not.
|
||||
This will enable gradients to be backpropagated through the features
|
||||
`latent_sde` in the code.
|
||||
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
|
||||
"""
|
||||
|
||||
def __init__(self, action_dim, full_std=True, use_expln=False,
|
||||
squash_output=False, epsilon=1e-6):
|
||||
squash_output=False, learn_features=False, epsilon=1e-6):
|
||||
super(StateDependentNoiseDistribution, self).__init__()
|
||||
self.distribution = None
|
||||
self.action_dim = action_dim
|
||||
self.latent_dim = None
|
||||
self.latent_sde_dim = None
|
||||
self.mean_actions = None
|
||||
self.log_std = None
|
||||
self.weights_dist = None
|
||||
self.exploration_mat = None
|
||||
self.exploration_matrices = None
|
||||
self.use_expln = use_expln
|
||||
self.full_std = full_std
|
||||
self.epsilon = epsilon
|
||||
self.learn_features = learn_features
|
||||
if squash_output:
|
||||
print("== Using TanhBijector ===")
|
||||
self.bijector = TanhBijector(epsilon)
|
||||
else:
|
||||
self.bijector = None
|
||||
|
|
@ -275,10 +272,11 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
if self.use_expln:
|
||||
# From SDE paper, it allows to keep variance
|
||||
# above zero and prevent it from growing too fast
|
||||
if log_std <= 0:
|
||||
std = th.exp(log_std)
|
||||
else:
|
||||
std = th.log(log_std + 1.0) + 1.0
|
||||
below_threshold = th.exp(log_std) * (log_std <= 0)
|
||||
# Avoid NaN: zeros values that are below zero
|
||||
safe_log_std = log_std * (log_std > 0) + self.epsilon
|
||||
above_threshold = (th.log1p(safe_log_std) + 1.0) * (log_std > 0)
|
||||
std = below_threshold + above_threshold
|
||||
else:
|
||||
# Use normal exponential
|
||||
std = th.exp(log_std)
|
||||
|
|
@ -286,57 +284,65 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
if self.full_std:
|
||||
return std
|
||||
# Reduce the number of parameters:
|
||||
return th.ones((self.latent_dim, self.action_dim)).to(log_std.device) * std
|
||||
return th.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std
|
||||
|
||||
def sample_weights(self, log_std):
|
||||
def sample_weights(self, log_std, batch_size=1):
|
||||
"""
|
||||
Sample weights for the noise exploration matrix,
|
||||
using a centered gaussian distribution.
|
||||
|
||||
:param log_std: (th.Tensor)
|
||||
:param batch_size: (int)
|
||||
"""
|
||||
std = self.get_std(log_std)
|
||||
self.weights_dist = Normal(th.zeros_like(std), std)
|
||||
self.exploration_mat = self.weights_dist.rsample()
|
||||
self.exploration_matrices = self.weights_dist.rsample((batch_size,))
|
||||
|
||||
def proba_distribution_net(self, latent_dim, log_std_init=-2.0):
|
||||
def proba_distribution_net(self, latent_dim, log_std_init=-2.0, latent_sde_dim=None):
|
||||
"""
|
||||
Create the layers and parameter that represent the distribution:
|
||||
one output will be the deterministic action, the other parameter will be the
|
||||
standard deviation of the distribution that control the weights of the noise matrix.
|
||||
|
||||
:param latent_dim: (int) Dimension og the last layer of the policy (before the action layer)
|
||||
:param latent_dim: (int) Dimension of the last layer of the policy (before the action layer)
|
||||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
:param latent_sde_dim: (int) Dimension of the last layer of the feature extractor
|
||||
for SDE. By default, it is shared with the policy network.
|
||||
:return: (nn.Linear, nn.Parameter)
|
||||
"""
|
||||
# Network for the deterministic action, it represents the mean of the distribution
|
||||
mean_actions_net = nn.Linear(latent_dim, self.action_dim)
|
||||
|
||||
self.latent_dim = latent_dim
|
||||
# When we learn features for the noise, the feature dimension
|
||||
# can be different between the policy and the noise network
|
||||
self.latent_sde_dim = latent_dim if latent_sde_dim is None else latent_sde_dim
|
||||
# Reduce the number of parameters if needed
|
||||
log_std = th.ones(latent_dim, self.action_dim) if self.full_std else th.ones(latent_dim, 1)
|
||||
log_std = th.ones(self.latent_sde_dim, self.action_dim) if self.full_std else th.ones(self.latent_sde_dim, 1)
|
||||
# Transform it to a parameter so it can be optimized
|
||||
log_std = nn.Parameter(log_std * log_std_init)
|
||||
# Sample an exploration matrix
|
||||
self.sample_weights(log_std)
|
||||
return mean_actions_net, log_std
|
||||
|
||||
def proba_distribution(self, mean_actions, log_std, latent_pi, deterministic=False):
|
||||
def proba_distribution(self, mean_actions, log_std, latent_sde, deterministic=False):
|
||||
"""
|
||||
Create and sample for the distribution given its parameters (mean, std)
|
||||
|
||||
:param mean_actions: (th.Tensor)
|
||||
:param log_std: (th.Tensor)
|
||||
:param latent_sde: (th.Tensor)
|
||||
:param deterministic: (bool)
|
||||
:return: (th.Tensor)
|
||||
"""
|
||||
variance = th.mm(latent_pi.detach() ** 2, self.get_std(log_std) ** 2)
|
||||
# Stop gradient if we don't want to influence the features
|
||||
latent_sde = latent_sde if self.learn_features else latent_sde.detach()
|
||||
variance = th.mm(latent_sde ** 2, self.get_std(log_std) ** 2)
|
||||
self.distribution = Normal(mean_actions, th.sqrt(variance + self.epsilon))
|
||||
|
||||
if deterministic:
|
||||
action = self.mode()
|
||||
else:
|
||||
action = self.sample(latent_pi)
|
||||
action = self.sample(latent_sde)
|
||||
return action, self
|
||||
|
||||
def mode(self):
|
||||
|
|
@ -345,11 +351,20 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
return self.bijector.forward(action)
|
||||
return action
|
||||
|
||||
def get_noise(self, latent_pi):
|
||||
return th.mm(latent_pi.detach(), self.exploration_mat)
|
||||
def get_noise(self, latent_sde):
|
||||
latent_sde = latent_sde if self.learn_features else latent_sde.detach()
|
||||
# Default case: only one exploration matrix
|
||||
if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices):
|
||||
return th.mm(latent_sde, self.exploration_mat)
|
||||
# Use batch matrix multiplication for efficient computation
|
||||
# (batch_size, n_features) -> (batch_size, 1, n_features)
|
||||
latent_sde = latent_sde.unsqueeze(1)
|
||||
# (batch_size, 1, n_actions)
|
||||
noise = th.bmm(latent_sde, self.exploration_matrices)
|
||||
return noise.squeeze(1)
|
||||
|
||||
def sample(self, latent_pi):
|
||||
noise = self.get_noise(latent_pi)
|
||||
def sample(self, latent_sde):
|
||||
noise = self.get_noise(latent_sde)
|
||||
action = self.distribution.mean + noise
|
||||
if self.bijector is not None:
|
||||
return self.bijector.forward(action)
|
||||
|
|
@ -359,8 +374,8 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
# TODO: account for the squashing?
|
||||
return self.distribution.entropy()
|
||||
|
||||
def log_prob_from_params(self, mean_actions, log_std, latent_pi):
|
||||
action, _ = self.proba_distribution(mean_actions, log_std, latent_pi)
|
||||
def log_prob_from_params(self, mean_actions, log_std, latent_sde):
|
||||
action, _ = self.proba_distribution(mean_actions, log_std, latent_sde)
|
||||
log_prob = self.log_prob(action)
|
||||
return action, log_prob
|
||||
|
||||
|
|
@ -391,11 +406,13 @@ class TanhBijector(object):
|
|||
|
||||
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
|
||||
"""
|
||||
|
||||
def __init__(self, epsilon=1e-6):
|
||||
super(TanhBijector, self).__init__()
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, x):
|
||||
@staticmethod
|
||||
def forward(x):
|
||||
return th.tanh(x)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -422,7 +439,7 @@ class TanhBijector(object):
|
|||
|
||||
def log_prob_correction(self, x):
|
||||
# Squash correction (from original SAC implementation)
|
||||
return th.log(1 - th.tanh(x) ** 2 + self.epsilon)
|
||||
return th.log(1.0 - th.tanh(x) ** 2 + self.epsilon)
|
||||
|
||||
|
||||
def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None):
|
||||
|
|
|
|||
|
|
@ -1,10 +1,34 @@
|
|||
# Copied from stable_baselines
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.vec_env import VecEnv
|
||||
|
||||
def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, render=False):
|
||||
|
||||
def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True,
|
||||
render=False, callback=None, reward_threshold=None,
|
||||
return_episode_rewards=False):
|
||||
"""
|
||||
Runs policy for n episodes and returns average reward
|
||||
Runs policy for `n_eval_episodes` episodes and returns average reward.
|
||||
This is made to work only with one env.
|
||||
|
||||
:param model: (BaseRLModel) The RL agent you want to evaluate.
|
||||
:param env: (gym.Env or VecEnv) The gym environment. In the case of a `VecEnv`
|
||||
this must contain only one environment.
|
||||
:param n_eval_episodes: (int) Number of episode to evaluate the agent
|
||||
:param deterministic: (bool) Whether to use deterministic or stochastic actions
|
||||
:param render: (bool) Whether to render the environment or not
|
||||
:param callback: (callable) callback function to do additional checks,
|
||||
called after each step.
|
||||
:param reward_threshold: (float) Minimum expected reward per episode,
|
||||
this will raise an error if the performance is not met
|
||||
:param return_episode_rewards: (bool) If True, a list of reward per episode
|
||||
will be returned instead of the mean.
|
||||
:return: (float, float) Mean reward per episode, std of reward per episode
|
||||
returns ([float], int) when `return_episode_rewards` is True
|
||||
"""
|
||||
if isinstance(env, VecEnv):
|
||||
assert env.num_envs == 1, "You must pass only one environment when using this function"
|
||||
|
||||
episode_rewards, n_steps = [], 0
|
||||
for _ in range(n_eval_episodes):
|
||||
obs = env.reset()
|
||||
|
|
@ -12,11 +36,19 @@ def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, render=F
|
|||
episode_reward = 0.0
|
||||
while not done:
|
||||
action = model.predict(obs, deterministic=deterministic)
|
||||
obs, reward, done, _ = env.step(action)
|
||||
obs, reward, done, _info = env.step(action)
|
||||
episode_reward += reward
|
||||
if callback is not None:
|
||||
callback(locals(), globals())
|
||||
n_steps += 1
|
||||
if render:
|
||||
env.render()
|
||||
episode_rewards.append(episode_reward)
|
||||
|
||||
return np.mean(episode_rewards), n_steps
|
||||
mean_reward = np.mean(episode_rewards)
|
||||
std_reward = np.std(episode_rewards)
|
||||
if reward_threshold is not None:
|
||||
assert mean_reward > reward_threshold, 'Mean reward below threshold: '\
|
||||
'{:.2f} < {:.2f}'.format(mean_reward, reward_threshold)
|
||||
if return_episode_rewards:
|
||||
return episode_rewards, n_steps
|
||||
return mean_reward, std_reward
|
||||
|
|
|
|||
|
|
@ -1,11 +1,10 @@
|
|||
"""
|
||||
Taken from stable-baselines
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import time
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import warnings
|
||||
from collections import defaultdict
|
||||
|
|
@ -185,15 +184,6 @@ class CSVOutputFormat(KVWriter):
|
|||
self.file.close()
|
||||
|
||||
|
||||
def summary_val(key, value):
|
||||
"""
|
||||
:param key: (str)
|
||||
:param value: (float)
|
||||
"""
|
||||
kwargs = {'tag': key, 'simple_value': float(value)}
|
||||
return tf.Summary.Value(**kwargs)
|
||||
|
||||
|
||||
def valid_float_value(value):
|
||||
"""
|
||||
Returns True if the value can be successfully cast into a float
|
||||
|
|
@ -515,3 +505,68 @@ def configure(folder=None, format_strs=None):
|
|||
|
||||
Logger.CURRENT = Logger(folder=folder, output_formats=output_formats)
|
||||
log('Logging to %s' % folder)
|
||||
|
||||
|
||||
def reset():
|
||||
"""
|
||||
reset the current logger
|
||||
"""
|
||||
if Logger.CURRENT is not Logger.DEFAULT:
|
||||
Logger.CURRENT.close()
|
||||
Logger.CURRENT = Logger.DEFAULT
|
||||
log('Reset logger')
|
||||
|
||||
|
||||
class ScopedConfigure(object):
|
||||
def __init__(self, folder=None, format_strs=None):
|
||||
"""
|
||||
Class for using context manager while logging
|
||||
|
||||
usage:
|
||||
with ScopedConfigure(folder=None, format_strs=None):
|
||||
{code}
|
||||
|
||||
:param folder: (str) the logging folder
|
||||
:param format_strs: ([str]) the list of output logging format
|
||||
"""
|
||||
self.dir = folder
|
||||
self.format_strs = format_strs
|
||||
self.prevlogger = None
|
||||
|
||||
def __enter__(self):
|
||||
self.prevlogger = Logger.CURRENT
|
||||
configure(folder=self.dir, format_strs=self.format_strs)
|
||||
|
||||
def __exit__(self, *args):
|
||||
Logger.CURRENT.close()
|
||||
Logger.CURRENT = self.prevlogger
|
||||
|
||||
|
||||
# ================================================================
|
||||
# Readers
|
||||
# ================================================================
|
||||
|
||||
def read_json(fname):
|
||||
"""
|
||||
read a json file using pandas
|
||||
|
||||
:param fname: (str) the file path to read
|
||||
:return: (pandas DataFrame) the data in the json
|
||||
"""
|
||||
import pandas
|
||||
data = []
|
||||
with open(fname, 'rt') as file_handler:
|
||||
for line in file_handler:
|
||||
data.append(json.loads(line))
|
||||
return pandas.DataFrame(data)
|
||||
|
||||
|
||||
def read_csv(fname):
|
||||
"""
|
||||
read a csv file using pandas
|
||||
|
||||
:param fname: (str) the file path to read
|
||||
:return: (pandas DataFrame) the data in the csv
|
||||
"""
|
||||
import pandas
|
||||
return pandas.read_csv(fname, index_col=None, comment='#')
|
||||
|
|
|
|||
|
|
@ -62,7 +62,25 @@ class BasePolicy(nn.Module):
|
|||
|
||||
def create_mlp(input_dim, output_dim, net_arch,
|
||||
activation_fn=nn.ReLU, squash_out=False):
|
||||
modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
|
||||
"""
|
||||
Create a multi layer perceptron (MLP), which is
|
||||
a collection of fully-connected layers each followed by an activation function.
|
||||
|
||||
:param input_dim: (int) Dimension of the input vector
|
||||
:param output_dim: (int)
|
||||
:param net_arch: ([int]) Architecture of the neural net
|
||||
It represents the number of units per layer.
|
||||
The length of this list is the number of layers.
|
||||
:param activation_fn: (th.nn.Module) The activation function
|
||||
to use after each layer.
|
||||
:param squash_out: (bool) Whether to squash the output using a Tanh
|
||||
activation function
|
||||
"""
|
||||
|
||||
if len(net_arch) > 0:
|
||||
modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
|
||||
else:
|
||||
modules = []
|
||||
|
||||
for idx in range(len(net_arch) - 1):
|
||||
modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))
|
||||
|
|
@ -75,9 +93,30 @@ def create_mlp(input_dim, output_dim, net_arch,
|
|||
return modules
|
||||
|
||||
|
||||
class BaseNetwork(nn.Module):
|
||||
"""docstring for BaseNetwork."""
|
||||
def create_sde_feature_extractor(features_dim, sde_net_arch, activation_fn):
|
||||
"""
|
||||
Create the neural network that will be used to extract features
|
||||
for the SDE.
|
||||
|
||||
:param features_dim: (int)
|
||||
:param sde_net_arch: ([int])
|
||||
:param activation_fn: (nn.Module)
|
||||
:return: (nn.Sequential, int)
|
||||
"""
|
||||
# Special case: when using states as features (i.e. sde_net_arch is an empty list)
|
||||
# don't use any activation function
|
||||
sde_activation = activation_fn if len(sde_net_arch) > 0 else None
|
||||
latent_sde_net = create_mlp(features_dim, -1, sde_net_arch, activation_fn=sde_activation, squash_out=False)
|
||||
latent_sde_dim = sde_net_arch[-1] if len(sde_net_arch) > 0 else features_dim
|
||||
sde_feature_extractor = nn.Sequential(*latent_sde_net)
|
||||
return sde_feature_extractor, latent_sde_dim
|
||||
|
||||
|
||||
class BaseNetwork(nn.Module):
|
||||
"""
|
||||
Abstract class for the different networks (actor/critic)
|
||||
that implements two helpers for using CEM with their weights.
|
||||
"""
|
||||
def __init__(self):
|
||||
super(BaseNetwork, self).__init__()
|
||||
|
||||
|
|
|
|||
|
|
@ -131,4 +131,4 @@ def json_to_data(json_string, custom_objects=None):
|
|||
else:
|
||||
# Read as it is
|
||||
return_data[data_key] = data_item
|
||||
return return_data
|
||||
return return_data
|
||||
|
|
|
|||
|
|
@ -35,4 +35,4 @@ def sync_envs_normalization(env, eval_env):
|
|||
if isinstance(env_tmp, VecNormalize):
|
||||
eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms)
|
||||
env_tmp = env_tmp.venv
|
||||
eval_env_tmp.venv
|
||||
eval_env_tmp = eval_env_tmp.venv
|
||||
|
|
|
|||
|
|
@ -58,7 +58,7 @@ class VecEnv(object):
|
|||
|
||||
:return: ([int] or [float]) observation
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def step_async(self, actions):
|
||||
|
|
@ -70,7 +70,7 @@ class VecEnv(object):
|
|||
You should not call this if a step_async run is
|
||||
already pending.
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def step_wait(self):
|
||||
|
|
@ -79,14 +79,14 @@ class VecEnv(object):
|
|||
|
||||
:return: ([int] or [float], [float], [bool], dict) observation, reward, done, information
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def close(self):
|
||||
"""
|
||||
Clean up the environment's resources.
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def get_attr(self, attr_name, indices=None):
|
||||
|
|
@ -97,7 +97,7 @@ class VecEnv(object):
|
|||
:param indices: (list,int) Indices of envs to get attribute from
|
||||
:return: (list) List of values of 'attr_name' in all environments
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def set_attr(self, attr_name, value, indices=None):
|
||||
|
|
@ -109,7 +109,7 @@ class VecEnv(object):
|
|||
:param indices: (list,int) Indices of envs to assign value
|
||||
:return: (NoneType)
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
@abstractmethod
|
||||
def env_method(self, method_name, *method_args, **method_kwargs):
|
||||
|
|
@ -122,7 +122,7 @@ class VecEnv(object):
|
|||
:param method_kwargs: (dict) Any keyword arguments to provide in the call
|
||||
:return: (list) List of items returned by the environment's method call
|
||||
"""
|
||||
pass
|
||||
raise NotImplementedError()
|
||||
|
||||
def step(self, actions):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -188,6 +188,17 @@ class SubprocVecEnv(VecEnv):
|
|||
for remote in target_remotes:
|
||||
remote.recv()
|
||||
|
||||
def seed(self, seed, indices=None):
|
||||
"""
|
||||
:param seed: (int or [int])
|
||||
:param indices: ([int])
|
||||
"""
|
||||
indices = self._get_indices(indices)
|
||||
if not hasattr(seed, 'len'):
|
||||
seed = [seed] * len(indices)
|
||||
assert len(seed) == len(indices)
|
||||
return [self.env_method('seed', seed[i], indices=i) for i in indices]
|
||||
|
||||
def env_method(self, method_name, *method_args, **method_kwargs):
|
||||
"""Call instance methods of vectorized environments."""
|
||||
indices = method_kwargs.get('indices')
|
||||
|
|
|
|||
|
|
@ -38,6 +38,45 @@ class VecNormalize(VecEnvWrapper):
|
|||
self.old_obs = np.array([])
|
||||
self.old_reward = np.array([])
|
||||
|
||||
def __getstate__(self):
|
||||
"""
|
||||
Gets state for pickling.
|
||||
|
||||
Excludes self.venv, as in general VecEnv's may not be pickleable."""
|
||||
state = self.__dict__.copy()
|
||||
# these attributes are not pickleable
|
||||
del state['venv']
|
||||
del state['class_attributes']
|
||||
# these attributes depend on the above and so we would prefer not to pickle
|
||||
del state['ret']
|
||||
return state
|
||||
|
||||
def __setstate__(self, state):
|
||||
"""
|
||||
Restores pickled state.
|
||||
|
||||
User must call set_venv() after unpickling before using.
|
||||
|
||||
:param state: (dict)"""
|
||||
self.__dict__.update(state)
|
||||
assert 'venv' not in state
|
||||
self.venv = None
|
||||
|
||||
def set_venv(self, venv):
|
||||
"""
|
||||
Sets the vector environment to wrap to venv.
|
||||
|
||||
Also sets attributes derived from this such as `num_env`.
|
||||
|
||||
:param venv: (VecEnv)
|
||||
"""
|
||||
if self.venv is not None:
|
||||
raise ValueError("Trying to set venv of already initialized VecNormalize wrapper.")
|
||||
VecEnvWrapper.__init__(self, venv)
|
||||
if self.obs_rms.mean.shape != self.observation_space.shape:
|
||||
raise ValueError("venv is incompatible with current statistics.")
|
||||
self.ret = np.zeros(self.num_envs)
|
||||
|
||||
def step_wait(self):
|
||||
"""
|
||||
Apply sequence of actions to sequence of environments
|
||||
|
|
@ -46,37 +85,44 @@ class VecNormalize(VecEnvWrapper):
|
|||
where 'news' is a boolean vector indicating whether each element is new.
|
||||
"""
|
||||
obs, rews, news, infos = self.venv.step_wait()
|
||||
self.ret = self.ret * self.gamma + rews
|
||||
self.old_obs = obs.copy()
|
||||
self.old_reward = rews.copy()
|
||||
obs = self._normalize_observation(obs)
|
||||
if self.norm_reward:
|
||||
if self.training:
|
||||
self.ret_rms.update(self.ret)
|
||||
rews = self.normalize_reward(rews)
|
||||
self.old_obs = obs
|
||||
self.old_rews = rews
|
||||
|
||||
if self.training:
|
||||
self.obs_rms.update(obs)
|
||||
obs = self.normalize_obs(obs)
|
||||
|
||||
if self.training:
|
||||
self._update_reward(rews)
|
||||
rews = self.normalize_reward(rews)
|
||||
|
||||
self.ret[news] = 0
|
||||
return obs, rews, news, infos
|
||||
|
||||
def _normalize_observation(self, obs):
|
||||
"""
|
||||
:param obs: (numpy tensor)
|
||||
"""
|
||||
if self.norm_obs:
|
||||
if self.training:
|
||||
self.obs_rms.update(obs)
|
||||
return self.normalize_obs(obs)
|
||||
else:
|
||||
return obs
|
||||
def _update_reward(self, reward):
|
||||
"""Update reward normalization statistics."""
|
||||
self.ret = self.ret * self.gamma + reward
|
||||
self.ret_rms.update(self.ret)
|
||||
|
||||
def normalize_obs(self, obs):
|
||||
"""
|
||||
Normalize observations using this VecNormalize's observations statistics.
|
||||
Calling this method does not update statistics.
|
||||
"""
|
||||
if self.norm_obs:
|
||||
return np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs,
|
||||
obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon),
|
||||
-self.clip_obs,
|
||||
self.clip_obs)
|
||||
return obs
|
||||
|
||||
def normalize_reward(self, reward):
|
||||
"""
|
||||
Normalize rewards using this VecNormalize's rewards statistics.
|
||||
Calling this method does not update statistics.
|
||||
"""
|
||||
if self.norm_reward:
|
||||
return np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
|
||||
reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon),
|
||||
-self.clip_reward, self.clip_reward)
|
||||
return reward
|
||||
|
||||
def unnormalize_obs(self, obs):
|
||||
|
|
@ -91,31 +137,45 @@ class VecNormalize(VecEnvWrapper):
|
|||
|
||||
def get_original_obs(self):
|
||||
"""
|
||||
returns the unnormalized observation
|
||||
|
||||
:return: (numpy float)
|
||||
Returns an unnormalized version of the observations from the most recent
|
||||
step or reset.
|
||||
"""
|
||||
return self.old_obs
|
||||
return self.old_obs.copy()
|
||||
|
||||
def get_original_reward(self):
|
||||
"""
|
||||
returns the unnormalized observation
|
||||
|
||||
:return: (numpy float)
|
||||
Returns an unnormalized version of the rewards from the most recent step.
|
||||
"""
|
||||
return self.old_reward
|
||||
return self.old_rews.copy()
|
||||
|
||||
def reset(self):
|
||||
"""
|
||||
Reset all environments
|
||||
"""
|
||||
obs = self.venv.reset()
|
||||
if len(np.array(obs).shape) == 1: # for when num_cpu is 1
|
||||
self.old_obs = [obs]
|
||||
else:
|
||||
self.old_obs = obs
|
||||
self.old_obs = obs
|
||||
self.ret = np.zeros(self.num_envs)
|
||||
return self._normalize_observation(obs)
|
||||
if self.training:
|
||||
self._update_reward(self.ret)
|
||||
return self.normalize_obs(obs)
|
||||
|
||||
@staticmethod
|
||||
def load(load_path, venv):
|
||||
"""
|
||||
Loads a saved VecNormalize object.
|
||||
|
||||
:param load_path: the path to load from.
|
||||
:param venv: the VecEnv to wrap.
|
||||
:return: (VecNormalize)
|
||||
"""
|
||||
with open(load_path, "rb") as file_handler:
|
||||
vec_normalize = pickle.load(file_handler)
|
||||
vec_normalize.set_venv(venv)
|
||||
return vec_normalize
|
||||
|
||||
def save(self, save_path):
|
||||
with open(save_path, "wb") as file_handler:
|
||||
pickle.dump(self, file_handler)
|
||||
|
||||
def save_running_average(self, path):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ import torch as th
|
|||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor
|
||||
from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor, \
|
||||
create_sde_feature_extractor
|
||||
from torchy_baselines.common.distributions import make_proba_distribution,\
|
||||
DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
|
||||
|
||||
|
|
@ -14,7 +15,7 @@ class PPOPolicy(BasePolicy):
|
|||
Policy class (with both actor and critic) for A2C and derivates (PPO).
|
||||
|
||||
:param observation_space: (gym.spaces.Space) Observation space
|
||||
:param action_dim: (gym.spaces.Space) Action space
|
||||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param learning_rate: (callable) Learning rate schedule (could be constant)
|
||||
:param net_arch: ([int or dict]) The specification of the policy and value networks.
|
||||
:param device: (str or th.device) Device on which the code should run.
|
||||
|
|
@ -25,16 +26,24 @@ class PPOPolicy(BasePolicy):
|
|||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
:param full_std: (bool) Whether to use (n_features x n_actions) parameters
|
||||
for the std instead of only (n_features,) when using SDE
|
||||
:param sde_net_arch: ([int]) Network architecture for extracting features
|
||||
when using SDE. If None, the latent features from the policy will be used.
|
||||
Pass an empty list to use the states as features.
|
||||
:param use_expln: (bool) Use `expln()` function instead of `exp()` to ensure
|
||||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
|
||||
:param squash_output: (bool) Whether to squash the output using a tanh function,
|
||||
this allows to ensure boundaries when using SDE.
|
||||
"""
|
||||
def __init__(self, observation_space, action_space,
|
||||
learning_rate, net_arch=None, device='cpu',
|
||||
activation_fn=nn.Tanh, adam_epsilon=1e-5,
|
||||
ortho_init=True, use_sde=False,
|
||||
log_std_init=0.0, full_std=True):
|
||||
log_std_init=0.0, full_std=True,
|
||||
sde_net_arch=None, use_expln=False, squash_output=False):
|
||||
super(PPOPolicy, self).__init__(observation_space, action_space, device)
|
||||
self.obs_dim = self.observation_space.shape[0]
|
||||
|
||||
|
||||
# Default network architecture, from stable-baselines
|
||||
if net_arch is None:
|
||||
net_arch = [dict(pi=[64, 64], vf=[64, 64])]
|
||||
|
|
@ -60,30 +69,49 @@ class PPOPolicy(BasePolicy):
|
|||
if use_sde:
|
||||
dist_kwargs = {
|
||||
'full_std': full_std,
|
||||
'squash_output': False,
|
||||
'use_expln': False
|
||||
'squash_output': squash_output,
|
||||
'use_expln': use_expln,
|
||||
'learn_features': sde_net_arch is not None
|
||||
}
|
||||
|
||||
self.sde_feature_extractor = None
|
||||
self.sde_net_arch = sde_net_arch
|
||||
|
||||
# Action distribution
|
||||
self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
|
||||
|
||||
self._build(learning_rate)
|
||||
|
||||
def reset_noise_net(self):
|
||||
def reset_noise(self, n_envs=1):
|
||||
"""
|
||||
Sample new weights for the exploration matrix.
|
||||
|
||||
:param n_envs: (int)
|
||||
"""
|
||||
self.action_dist.sample_weights(self.log_std)
|
||||
self.action_dist.sample_weights(self.log_std, batch_size=n_envs)
|
||||
|
||||
def _build(self, learning_rate):
|
||||
self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch,
|
||||
activation_fn=self.activation_fn, device=self.device)
|
||||
|
||||
if isinstance(self.action_dist, (DiagGaussianDistribution, StateDependentNoiseDistribution)):
|
||||
self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
|
||||
latent_dim_pi = self.mlp_extractor.latent_dim_pi
|
||||
|
||||
# Separate feature extractor for SDE
|
||||
if self.sde_net_arch is not None:
|
||||
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(self.features_dim,
|
||||
self.sde_net_arch,
|
||||
self.activation_fn)
|
||||
|
||||
if isinstance(self.action_dist, DiagGaussianDistribution):
|
||||
self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi,
|
||||
log_std_init=self.log_std_init)
|
||||
elif isinstance(self.action_dist, StateDependentNoiseDistribution):
|
||||
latent_sde_dim = latent_dim_pi if self.sde_net_arch is None else latent_sde_dim
|
||||
self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi,
|
||||
latent_sde_dim=latent_sde_dim,
|
||||
log_std_init=self.log_std_init)
|
||||
elif isinstance(self.action_dist, CategoricalDistribution):
|
||||
self.action_net = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi)
|
||||
self.action_net = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi)
|
||||
|
||||
self.value_net = nn.Linear(self.mlp_extractor.latent_dim_vf, 1)
|
||||
# Init weights: use orthogonal initialization
|
||||
|
|
@ -102,30 +130,39 @@ class PPOPolicy(BasePolicy):
|
|||
def forward(self, obs, deterministic=False):
|
||||
if not isinstance(obs, th.Tensor):
|
||||
obs = th.FloatTensor(obs).to(self.device)
|
||||
latent_pi, latent_vf = self._get_latent(obs)
|
||||
latent_pi, latent_vf, latent_sde = self._get_latent(obs)
|
||||
value = self.value_net(latent_vf)
|
||||
action, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
|
||||
action, action_distribution = self._get_action_dist_from_latent(latent_pi, latent_sde=latent_sde,
|
||||
deterministic=deterministic)
|
||||
log_prob = action_distribution.log_prob(action)
|
||||
return action, value, log_prob
|
||||
|
||||
def _get_latent(self, obs):
|
||||
return self.mlp_extractor(self.features_extractor(obs))
|
||||
features = self.features_extractor(obs)
|
||||
latent_pi, latent_vf = self.mlp_extractor(features)
|
||||
# Features for sde
|
||||
latent_sde = latent_pi
|
||||
if self.sde_feature_extractor is not None:
|
||||
latent_sde = self.sde_feature_extractor(features)
|
||||
return latent_pi, latent_vf, latent_sde
|
||||
|
||||
def _get_action_dist_from_latent(self, latent_pi, deterministic=False):
|
||||
def _get_action_dist_from_latent(self, latent_pi, latent_sde=None, deterministic=False):
|
||||
mean_actions = self.action_net(latent_pi)
|
||||
|
||||
if isinstance(self.action_dist, DiagGaussianDistribution):
|
||||
return self.action_dist.proba_distribution(mean_actions, self.log_std, deterministic=deterministic)
|
||||
|
||||
elif isinstance(self.action_dist, CategoricalDistribution):
|
||||
# Here mean_actions are the logits before the softmax
|
||||
return self.action_dist.proba_distribution(mean_actions, deterministic=deterministic)
|
||||
|
||||
elif isinstance(self.action_dist, StateDependentNoiseDistribution):
|
||||
return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_pi, deterministic=deterministic)
|
||||
return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde,
|
||||
deterministic=deterministic)
|
||||
|
||||
def actor_forward(self, obs, deterministic=False):
|
||||
latent_pi, _ = self._get_latent(obs)
|
||||
action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
|
||||
latent_pi, _, latent_sde = self._get_latent(obs)
|
||||
action, _ = self._get_action_dist_from_latent(latent_pi, latent_sde, deterministic=deterministic)
|
||||
return action.detach().cpu().numpy()
|
||||
|
||||
def evaluate_actions(self, obs, action, deterministic=False):
|
||||
|
|
@ -139,14 +176,14 @@ class PPOPolicy(BasePolicy):
|
|||
:return: (th.Tensor, th.Tensor, th.Tensor) estimated value, log likelihood of taking those actions
|
||||
and entropy of the action distribution.
|
||||
"""
|
||||
latent_pi, latent_vf = self._get_latent(obs)
|
||||
_, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
|
||||
latent_pi, latent_vf, latent_sde = self._get_latent(obs)
|
||||
_, action_distribution = self._get_action_dist_from_latent(latent_pi, latent_sde, deterministic=deterministic)
|
||||
log_prob = action_distribution.log_prob(action)
|
||||
value = self.value_net(latent_vf)
|
||||
return value, log_prob, action_distribution.entropy()
|
||||
|
||||
def value_forward(self, obs):
|
||||
_, latent_vf = self._get_latent(obs)
|
||||
_, latent_vf, _ = self._get_latent(obs)
|
||||
return self.value_net(latent_vf)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -14,10 +14,8 @@ except ImportError:
|
|||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.base_class import BaseRLModel
|
||||
from torchy_baselines.common.evaluation import evaluate_policy
|
||||
from torchy_baselines.common.buffers import RolloutBuffer
|
||||
from torchy_baselines.common.utils import explained_variance, get_schedule_fn
|
||||
from torchy_baselines.common.vec_env import sync_envs_normalization
|
||||
from torchy_baselines.common import logger
|
||||
from torchy_baselines.ppo.policies import PPOPolicy
|
||||
|
||||
|
|
@ -43,7 +41,8 @@ class PPO(BaseRLModel):
|
|||
:param n_epochs: (int) Number of epoch when optimizing the surrogate loss
|
||||
:param gamma: (float) Discount factor
|
||||
:param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
|
||||
:param clip_range: (float or callable) Clipping parameter, it can be a function of the current progress (from 1 to 0).
|
||||
:param clip_range: (float or callable) Clipping parameter, it can be a function of the current progress
|
||||
(from 1 to 0).
|
||||
:param clip_range_vf: (float or callable) Clipping parameter for the value function,
|
||||
it can be a function of the current progress (from 1 to 0).
|
||||
This is a parameter specific to the OpenAI implementation. If None is passed (default),
|
||||
|
|
@ -54,6 +53,8 @@ class PPO(BaseRLModel):
|
|||
:param max_grad_norm: (float) The maximum value for the gradient clipping
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
|
||||
instead of action noise exploration (default: False)
|
||||
:param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
|
||||
Default: -1 (only sample at the beginning of the rollout)
|
||||
:param target_kl: (float) Limit the KL divergence between updates,
|
||||
because the clipping is not enough to prevent large update
|
||||
see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213)
|
||||
|
|
@ -68,17 +69,17 @@ class PPO(BaseRLModel):
|
|||
Setting it to auto, the code will be run on the GPU if possible.
|
||||
:param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance
|
||||
"""
|
||||
|
||||
def __init__(self, policy, env, learning_rate=3e-4,
|
||||
n_steps=2048, batch_size=64, n_epochs=10,
|
||||
gamma=0.99, gae_lambda=0.95, clip_range=0.2, clip_range_vf=None,
|
||||
ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5, use_sde=False,
|
||||
ent_coef=0.0, vf_coef=0.5, max_grad_norm=0.5,
|
||||
use_sde=False, sde_sample_freq=-1,
|
||||
target_kl=None, tensorboard_log=None, create_eval_env=False,
|
||||
policy_kwargs=None, verbose=0, seed=0, device='auto',
|
||||
_init_setup_model=True):
|
||||
|
||||
super(PPO, self).__init__(policy, env, PPOPolicy, policy_kwargs=policy_kwargs,
|
||||
verbose=verbose, device=device,
|
||||
verbose=verbose, device=device, use_sde=use_sde, sde_sample_freq=sde_sample_freq,
|
||||
create_eval_env=create_eval_env, support_multi_env=True, seed=seed)
|
||||
|
||||
self.learning_rate = learning_rate
|
||||
|
|
@ -96,7 +97,6 @@ class PPO(BaseRLModel):
|
|||
self.target_kl = target_kl
|
||||
self.tensorboard_log = tensorboard_log
|
||||
self.tb_writer = None
|
||||
self.use_sde = use_sde
|
||||
|
||||
if _init_setup_model:
|
||||
self._setup_model()
|
||||
|
|
@ -157,9 +157,13 @@ class PPO(BaseRLModel):
|
|||
# Sample new weights for the state dependent exploration
|
||||
# TODO: ensure episodic setting?
|
||||
if self.use_sde:
|
||||
self.policy.reset_noise_net()
|
||||
self.policy.reset_noise(env.num_envs)
|
||||
|
||||
while n_steps < n_rollout_steps:
|
||||
if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
|
||||
# Sample a new noise matrix
|
||||
self.policy.reset_noise(env.num_envs)
|
||||
|
||||
with th.no_grad():
|
||||
actions, values, log_probs = self.policy.forward(obs)
|
||||
actions = actions.cpu().numpy()
|
||||
|
|
@ -204,6 +208,12 @@ class PPO(BaseRLModel):
|
|||
# Convert discrete action for float to long
|
||||
action = action.long().flatten()
|
||||
|
||||
# Re-sample the noise matrix because the log_std has changed
|
||||
# TODO: investigate why there is no issue with the gradient
|
||||
# if that line is commented (as in SAC)
|
||||
if self.use_sde:
|
||||
self.policy.reset_noise(batch_size)
|
||||
|
||||
values, log_prob, entropy = self.policy.evaluate_actions(obs, action)
|
||||
values = values.flatten()
|
||||
# Normalize advantage
|
||||
|
|
@ -291,27 +301,20 @@ class PPO(BaseRLModel):
|
|||
|
||||
self.train(self.n_epochs, batch_size=self.batch_size)
|
||||
|
||||
# Evaluate agent
|
||||
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
|
||||
timesteps_since_eval %= eval_freq
|
||||
sync_envs_normalization(self.env, eval_env)
|
||||
|
||||
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
|
||||
if self.tb_writer is not None:
|
||||
self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
|
||||
|
||||
evaluations.append(mean_reward)
|
||||
if self.verbose > 0:
|
||||
print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
|
||||
print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
|
||||
# Evaluate the agent
|
||||
timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
|
||||
timesteps_since_eval, deterministic=True)
|
||||
# For tensorboard integration
|
||||
# if self.tb_writer is not None:
|
||||
# self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
|
||||
|
||||
return self
|
||||
|
||||
def get_opt_parameters(self):
|
||||
"""
|
||||
Returns a dict of all the optimizers and their parameters
|
||||
|
||||
:return: (dict) of optimizer names and their state_dict
|
||||
|
||||
:return: (dict) of optimizer names and their state_dict
|
||||
"""
|
||||
return {"opt": self.policy.optimizer.state_dict()}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,46 +1,160 @@
|
|||
import torch as th
|
||||
import torch.nn as nn
|
||||
|
||||
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork
|
||||
from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution
|
||||
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, \
|
||||
create_sde_feature_extractor
|
||||
from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
|
||||
|
||||
# CAP the standard deviation of the actor
|
||||
LOG_STD_MAX = 2
|
||||
LOG_STD_MIN = -20
|
||||
|
||||
|
||||
class LeakyClip(nn.Module):
|
||||
"""
|
||||
Cip values outside a certain range
|
||||
(it is not a hard clip, there is a small slope to have non-zero gradient)
|
||||
|
||||
:param min_val: (float)
|
||||
:param max_val: (float)
|
||||
:param slope: (float)
|
||||
"""
|
||||
def __init__(self, min_val=-2.0, max_val=2.0, slope=0.01):
|
||||
super(LeakyClip, self).__init__()
|
||||
self.min_val = min_val
|
||||
self.max_val = max_val
|
||||
self.slope = slope
|
||||
|
||||
def forward(self, x):
|
||||
linear_part = x * (x >= self.min_val) * (x <= self.max_val)
|
||||
above_max_val = self.slope * (x - self.max_val) * (x > self.max_val)
|
||||
below_min_val = self.slope * (x - self.min_val) * (x < self.min_val)
|
||||
return linear_part + below_min_val + above_max_val
|
||||
|
||||
|
||||
class Actor(BaseNetwork):
|
||||
def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU):
|
||||
"""
|
||||
Actor network (policy) for SAC.
|
||||
|
||||
:param obs_dim: (int) Dimension of the observation
|
||||
:param action_dim: (int) Dimension of the action space
|
||||
:param net_arch: ([int]) Network architecture
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration or not
|
||||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
:param full_std: (bool) Whether to use (n_features x n_actions) parameters
|
||||
for the std instead of only (n_features,) when using SDE.
|
||||
:param sde_net_arch: ([int]) Network architecture for extracting features
|
||||
when using SDE. If None, the latent features from the policy will be used.
|
||||
Pass an empty list to use the states as features.
|
||||
:param use_expln: (bool) Use `expln()` function instead of `exp()` when using SDE to ensure
|
||||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
|
||||
"""
|
||||
def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
|
||||
use_sde=False, log_std_init=-3, full_std=True,
|
||||
sde_net_arch=None, use_expln=False):
|
||||
super(Actor, self).__init__()
|
||||
|
||||
# TODO: orthogonal initialization?
|
||||
actor_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
|
||||
self.actor_net = nn.Sequential(*actor_net)
|
||||
latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
|
||||
self.latent_pi = nn.Sequential(*latent_pi_net)
|
||||
self.use_sde = use_sde
|
||||
self.sde_feature_extractor = None
|
||||
|
||||
self.action_dist = SquashedDiagGaussianDistribution(action_dim)
|
||||
self.mu = nn.Linear(net_arch[-1], action_dim)
|
||||
self.log_std = nn.Linear(net_arch[-1], action_dim)
|
||||
if self.use_sde:
|
||||
latent_sde_dim = net_arch[-1]
|
||||
# Separate feature extractor for SDE
|
||||
if sde_net_arch is not None:
|
||||
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch,
|
||||
activation_fn)
|
||||
|
||||
# TODO: check for the learn_features
|
||||
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln,
|
||||
learn_features=True, squash_output=True)
|
||||
self.mu, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
|
||||
latent_sde_dim=latent_sde_dim,
|
||||
log_std_init=log_std_init)
|
||||
# Avoid saturation by limiting the mean of the gaussian to be in [-1, 1]
|
||||
# self.mu = nn.Sequential(self.mu, nn.Tanh())
|
||||
self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-2.0, max_val=2.0))
|
||||
# Small positive slope to have non-zero gradient
|
||||
# self.mu = nn.Sequential(self.mu, LeakyClip())
|
||||
else:
|
||||
self.action_dist = SquashedDiagGaussianDistribution(action_dim)
|
||||
self.mu = nn.Linear(net_arch[-1], action_dim)
|
||||
self.log_std = nn.Linear(net_arch[-1], action_dim)
|
||||
|
||||
def get_std(self):
|
||||
"""
|
||||
Retrieve the standard deviation of the action distribution.
|
||||
Only useful when using SDE.
|
||||
It corresponds to `th.exp(log_std)` in the normal case,
|
||||
but is slightly different when using `expln` function
|
||||
(cf StateDependentNoiseDistribution doc).
|
||||
|
||||
:return: (th.Tensor)
|
||||
"""
|
||||
return self.action_dist.get_std(self.log_std)
|
||||
|
||||
def reset_noise(self, batch_size=1):
|
||||
"""
|
||||
Sample new weights for the exploration matrix, when using SDE.
|
||||
|
||||
:param batch_size: (int)
|
||||
"""
|
||||
self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
|
||||
|
||||
def _get_latent(self, obs):
|
||||
latent_pi = self.latent_pi(obs)
|
||||
|
||||
if self.sde_feature_extractor is not None:
|
||||
latent_sde = self.sde_feature_extractor(obs)
|
||||
else:
|
||||
latent_sde = latent_pi
|
||||
return latent_pi, latent_sde
|
||||
|
||||
def get_action_dist_params(self, obs):
|
||||
latent = self.actor_net(obs)
|
||||
mean_actions, log_std = self.mu(latent), self.log_std(latent)
|
||||
# Original Implementation to cap the standard deviation
|
||||
log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)
|
||||
return mean_actions, log_std
|
||||
latent_pi, latent_sde = self._get_latent(obs)
|
||||
|
||||
if self.use_sde:
|
||||
mean_actions, log_std = self.mu(latent_pi), self.log_std
|
||||
else:
|
||||
mean_actions, log_std = self.mu(latent_pi), self.log_std(latent_pi)
|
||||
# Original Implementation to cap the standard deviation
|
||||
log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)
|
||||
return mean_actions, log_std, latent_sde
|
||||
|
||||
def forward(self, obs, deterministic=False):
|
||||
mean_actions, log_std = self.get_action_dist_params(obs)
|
||||
# Note the action is squashed
|
||||
action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic)
|
||||
mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
|
||||
if self.use_sde:
|
||||
# Note the action is squashed
|
||||
action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent_sde,
|
||||
deterministic=deterministic)
|
||||
else:
|
||||
# Note the action is squashed
|
||||
action, _ = self.action_dist.proba_distribution(mean_actions, log_std,
|
||||
deterministic=deterministic)
|
||||
return action
|
||||
|
||||
def action_log_prob(self, obs):
|
||||
mean_actions, log_std = self.get_action_dist_params(obs)
|
||||
action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std)
|
||||
mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
|
||||
|
||||
if self.use_sde:
|
||||
action, log_prob = self.action_dist.log_prob_from_params(mean_actions, self.log_std, latent_sde)
|
||||
else:
|
||||
action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std)
|
||||
return action, log_prob
|
||||
|
||||
|
||||
class Critic(BaseNetwork):
|
||||
"""
|
||||
Critic network (q-value function) for SAC.
|
||||
|
||||
:param obs_dim: (int) Dimension of the observation
|
||||
:param action_dim: (int) Dimension of the action space
|
||||
:param net_arch: ([int]) Network architecture
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
"""
|
||||
def __init__(self, obs_dim, action_dim,
|
||||
net_arch, activation_fn=nn.ReLU):
|
||||
super(Critic, self).__init__()
|
||||
|
|
@ -62,9 +176,28 @@ class Critic(BaseNetwork):
|
|||
|
||||
|
||||
class SACPolicy(BasePolicy):
|
||||
"""
|
||||
Policy class (with both actor and critic) for SAC.
|
||||
|
||||
:param observation_space: (gym.spaces.Space) Observation space
|
||||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param learning_rate: (callable) Learning rate schedule (could be constant)
|
||||
:param net_arch: ([int or dict]) The specification of the policy and value networks.
|
||||
:param device: (str or th.device) Device on which the code should run.
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration or not
|
||||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
:param sde_net_arch: ([int]) Network architecture for extracting features
|
||||
when using SDE. If None, the latent features from the policy will be used.
|
||||
Pass an empty list to use the states as features.
|
||||
:param use_expln: (bool) Use `expln()` function instead of `exp()` when using SDE to ensure
|
||||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
|
||||
"""
|
||||
def __init__(self, observation_space, action_space,
|
||||
learning_rate, net_arch=None, device='cpu',
|
||||
activation_fn=nn.ReLU):
|
||||
activation_fn=nn.ReLU, use_sde=False,
|
||||
log_std_init=-3, sde_net_arch=None, use_expln=False):
|
||||
super(SACPolicy, self).__init__(observation_space, action_space, device)
|
||||
|
||||
if net_arch is None:
|
||||
|
|
@ -80,6 +213,14 @@ class SACPolicy(BasePolicy):
|
|||
'net_arch': self.net_arch,
|
||||
'activation_fn': self.activation_fn
|
||||
}
|
||||
self.actor_kwargs = self.net_args.copy()
|
||||
sde_kwargs = {
|
||||
'use_sde': use_sde,
|
||||
'log_std_init': log_std_init,
|
||||
'sde_net_arch': sde_net_arch,
|
||||
'use_expln': use_expln
|
||||
}
|
||||
self.actor_kwargs.update(sde_kwargs)
|
||||
self.actor, self.actor_target = None, None
|
||||
self.critic, self.critic_target = None, None
|
||||
|
||||
|
|
@ -95,11 +236,14 @@ class SACPolicy(BasePolicy):
|
|||
self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
|
||||
|
||||
def make_actor(self):
|
||||
return Actor(**self.net_args).to(self.device)
|
||||
return Actor(**self.actor_kwargs).to(self.device)
|
||||
|
||||
def make_critic(self):
|
||||
return Critic(**self.net_args).to(self.device)
|
||||
|
||||
def forward(self, obs):
|
||||
return self.actor(obs)
|
||||
|
||||
|
||||
MlpPolicy = SACPolicy
|
||||
|
||||
|
|
|
|||
|
|
@ -1,14 +1,11 @@
|
|||
import time
|
||||
|
||||
import torch as th
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.base_class import BaseRLModel
|
||||
from torchy_baselines.common.buffers import ReplayBuffer
|
||||
from torchy_baselines.common.evaluation import evaluate_policy
|
||||
from torchy_baselines.sac.policies import SACPolicy
|
||||
from torchy_baselines.common.vec_env import sync_envs_normalization
|
||||
from torchy_baselines.common import logger
|
||||
|
||||
|
||||
class SAC(BaseRLModel):
|
||||
|
|
@ -44,6 +41,10 @@ class SAC(BaseRLModel):
|
|||
:param action_noise: (ActionNoise) the action noise type (None by default), this can help
|
||||
for hard exploration problem. Cf common.noise for the different action noise type.
|
||||
:param gamma: (float) the discount factor
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
|
||||
instead of action noise exploration (default: False)
|
||||
:param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
|
||||
Default: -1 (only sample at the beginning of the rollout)
|
||||
:param create_eval_env: (bool) Whether to create a second environment that will be
|
||||
used for evaluating the agent periodically. (Only available when passing string for the environment)
|
||||
:param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
|
||||
|
|
@ -59,12 +60,14 @@ class SAC(BaseRLModel):
|
|||
tau=0.005, ent_coef='auto', target_update_interval=1,
|
||||
train_freq=1, gradient_steps=1, n_episodes_rollout=-1,
|
||||
target_entropy='auto', action_noise=None,
|
||||
gamma=0.99, tensorboard_log=None, create_eval_env=False,
|
||||
gamma=0.99, use_sde=False, sde_sample_freq=-1,
|
||||
tensorboard_log=None, create_eval_env=False,
|
||||
policy_kwargs=None, verbose=0, seed=0, device='auto',
|
||||
_init_setup_model=True):
|
||||
|
||||
super(SAC, self).__init__(policy, env, SACPolicy, policy_kwargs, verbose, device,
|
||||
create_eval_env=create_eval_env, seed=seed)
|
||||
create_eval_env=create_eval_env, seed=seed,
|
||||
use_sde=use_sde, sde_sample_freq=sde_sample_freq)
|
||||
|
||||
self.learning_rate = learning_rate
|
||||
self.target_entropy = target_entropy
|
||||
|
|
@ -85,6 +88,7 @@ class SAC(BaseRLModel):
|
|||
self.n_episodes_rollout = n_episodes_rollout
|
||||
self.action_noise = action_noise
|
||||
self.gamma = gamma
|
||||
self.ent_coef_optimizer = None
|
||||
|
||||
if _init_setup_model:
|
||||
self._setup_model()
|
||||
|
|
@ -122,11 +126,12 @@ class SAC(BaseRLModel):
|
|||
# Force conversion to float
|
||||
# this will throw an error if a malformed string (different from 'auto')
|
||||
# is passed
|
||||
self.ent_coef = float(self.ent_coef)
|
||||
self.ent_coef = th.tensor(float(self.ent_coef)).to(self.device)
|
||||
|
||||
self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
|
||||
self.policy = self.policy_class(self.observation_space, self.action_space,
|
||||
self.learning_rate, device=self.device, **self.policy_kwargs)
|
||||
self.learning_rate, use_sde=self.use_sde,
|
||||
device=self.device, **self.policy_kwargs)
|
||||
self.policy = self.policy.to(self.device)
|
||||
self._create_aliases()
|
||||
|
||||
|
|
@ -168,12 +173,21 @@ class SAC(BaseRLModel):
|
|||
|
||||
obs, action_batch, next_obs, done, reward = replay_data
|
||||
|
||||
# Two options: retain_graph=True in the actor_loss.backward()
|
||||
# or sample again the noise matrix
|
||||
# otherwise the intermediate step `std = th.exp(log_std)`
|
||||
# is lost and we cannot backpropagate through again
|
||||
# anyway, we need to sample because `log_std` may have changed between two gradient steps
|
||||
if self.use_sde:
|
||||
self.actor.reset_noise(batch_size=batch_size)
|
||||
# self.actor.reset_noise()
|
||||
|
||||
# Action by the current actor for the sampled state
|
||||
action_pi, log_prob = self.actor.action_log_prob(obs)
|
||||
log_prob = log_prob.reshape(-1, 1)
|
||||
|
||||
ent_coef_loss = None
|
||||
if not isinstance(self.ent_coef, float):
|
||||
if self.ent_coef_optimizer is not None:
|
||||
# Important: detach the variable from the graph
|
||||
# so we don't change it with other losses
|
||||
# see https://github.com/rail-berkeley/softlearning/issues/60
|
||||
|
|
@ -189,10 +203,11 @@ class SAC(BaseRLModel):
|
|||
ent_coef_loss.backward()
|
||||
self.ent_coef_optimizer.step()
|
||||
|
||||
# Select action according to policy
|
||||
next_action, next_log_prob = self.actor.action_log_prob(next_obs)
|
||||
|
||||
with th.no_grad():
|
||||
# if self.use_sde:
|
||||
# self.actor.reset_noise(batch_size=batch_size)
|
||||
# Select action according to policy
|
||||
next_action, next_log_prob = self.actor.action_log_prob(next_obs)
|
||||
# Compute the target Q value
|
||||
target_q1, target_q2 = self.critic_target(next_obs, next_action)
|
||||
target_q = th.min(target_q1, target_q2)
|
||||
|
|
@ -228,6 +243,13 @@ class SAC(BaseRLModel):
|
|||
for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
|
||||
target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
|
||||
|
||||
# TODO: average
|
||||
logger.logkv("ent_coef", ent_coef.item())
|
||||
logger.logkv("actor_loss", actor_loss.item())
|
||||
logger.logkv("critic_loss", critic_loss.item())
|
||||
if ent_coef_loss is not None:
|
||||
logger.logkv("ent_coef_loss", ent_coef_loss.item())
|
||||
|
||||
def learn(self, total_timesteps, callback=None, log_interval=4,
|
||||
eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="SAC",
|
||||
reset_num_timesteps=True):
|
||||
|
|
@ -262,15 +284,8 @@ class SAC(BaseRLModel):
|
|||
|
||||
self.train(gradient_steps, batch_size=self.batch_size)
|
||||
|
||||
# Evaluate episode
|
||||
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
|
||||
timesteps_since_eval %= eval_freq
|
||||
sync_envs_normalization(self.env, eval_env)
|
||||
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
|
||||
evaluations.append(mean_reward)
|
||||
if self.verbose > 0:
|
||||
print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
|
||||
print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
|
||||
timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
|
||||
timesteps_since_eval, deterministic=True)
|
||||
|
||||
return self
|
||||
|
||||
|
|
@ -278,7 +293,7 @@ class SAC(BaseRLModel):
|
|||
"""
|
||||
Returns a dict of all the optimizers and their parameters
|
||||
|
||||
:return: (Dict) of optimizer names and their state_dict
|
||||
:return: (Dict) of optimizer names and their state_dict
|
||||
"""
|
||||
opt_dict = {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict()}
|
||||
if self.ent_coef_optimizer is not None:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
import torch as th
|
||||
import torch.nn as nn
|
||||
|
||||
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork
|
||||
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, \
|
||||
create_sde_feature_extractor
|
||||
from torchy_baselines.common.distributions import StateDependentNoiseDistribution
|
||||
|
||||
|
||||
|
|
@ -19,10 +20,16 @@ class Actor(BaseNetwork):
|
|||
:param lr_sde: (float) Learning rate for the standard deviation of the noise
|
||||
:param full_std: (bool) Whether to use (n_features x n_actions) parameters
|
||||
for the std instead of only (n_features,) when using SDE.
|
||||
:param sde_net_arch: ([int]) Network architecture for extracting features
|
||||
when using SDE. If None, the latent features from the policy will be used.
|
||||
Pass an empty list to use the states as features.
|
||||
:param use_expln: (bool) Use `expln()` function instead of `exp()` when using SDE to ensure
|
||||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
|
||||
"""
|
||||
def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU,
|
||||
use_sde=False, log_std_init=-2, clip_noise=None,
|
||||
lr_sde=3e-4, full_std=False):
|
||||
use_sde=False, log_std_init=-3, clip_noise=None,
|
||||
lr_sde=3e-4, full_std=False, sde_net_arch=None, use_expln=False):
|
||||
super(Actor, self).__init__()
|
||||
|
||||
self.latent_pi, self.log_std = None, None
|
||||
|
|
@ -30,23 +37,33 @@ class Actor(BaseNetwork):
|
|||
self.use_sde, self.sde_optimizer = use_sde, None
|
||||
self.action_dim = action_dim
|
||||
self.full_std = full_std
|
||||
self.sde_feature_extractor = None
|
||||
|
||||
if use_sde:
|
||||
latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False)
|
||||
self.latent_pi = nn.Sequential(*latent_pi)
|
||||
latent_pi_net = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False)
|
||||
self.latent_pi = nn.Sequential(*latent_pi_net)
|
||||
latent_sde_dim = net_arch[-1]
|
||||
learn_features = sde_net_arch is not None
|
||||
|
||||
# Separate feature extractor for SDE
|
||||
if sde_net_arch is not None:
|
||||
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch,
|
||||
activation_fn)
|
||||
|
||||
# Create state dependent noise matrix (SDE)
|
||||
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
|
||||
squash_output=False)
|
||||
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=use_expln,
|
||||
squash_output=False, learn_features=learn_features)
|
||||
action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
|
||||
latent_sde_dim=latent_sde_dim,
|
||||
log_std_init=log_std_init)
|
||||
# Squash output
|
||||
self.actor_net = nn.Sequential(action_net, nn.Tanh())
|
||||
self.mu = nn.Sequential(action_net, nn.Tanh())
|
||||
self.clip_noise = clip_noise
|
||||
self.sde_optimizer = th.optim.Adam([self.log_std], lr=lr_sde)
|
||||
self.reset_noise()
|
||||
else:
|
||||
actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True)
|
||||
self.actor_net = nn.Sequential(*actor_net)
|
||||
self.mu = nn.Sequential(*actor_net)
|
||||
|
||||
def get_std(self):
|
||||
"""
|
||||
|
|
@ -60,9 +77,18 @@ class Actor(BaseNetwork):
|
|||
"""
|
||||
return self.action_dist.get_std(self.log_std)
|
||||
|
||||
def _get_action_dist_from_latent(self, latent_pi):
|
||||
mean_actions = self.actor_net(latent_pi)
|
||||
return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_pi)
|
||||
def _get_action_dist_from_latent(self, latent_pi, latent_sde):
|
||||
mean_actions = self.mu(latent_pi)
|
||||
return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde)
|
||||
|
||||
def _get_latent(self, obs):
|
||||
latent_pi = self.latent_pi(obs)
|
||||
|
||||
if self.sde_feature_extractor is not None:
|
||||
latent_sde = self.sde_feature_extractor(obs)
|
||||
else:
|
||||
latent_sde = latent_pi
|
||||
return latent_pi, latent_sde
|
||||
|
||||
def evaluate_actions(self, obs, action):
|
||||
"""
|
||||
|
|
@ -71,39 +97,38 @@ class Actor(BaseNetwork):
|
|||
|
||||
:param obs: (th.Tensor)
|
||||
:param action: (th.Tensor)
|
||||
:param deterministic: (bool)
|
||||
:return: (th.Tensor, th.Tensor) log likelihood of taking those actions
|
||||
and entropy of the action distribution.
|
||||
"""
|
||||
with th.no_grad():
|
||||
latent_pi = self.latent_pi(obs)
|
||||
_, distribution = self._get_action_dist_from_latent(latent_pi)
|
||||
latent_pi, latent_sde = self._get_latent(obs)
|
||||
_, distribution = self._get_action_dist_from_latent(latent_pi, latent_sde)
|
||||
log_prob = distribution.log_prob(action)
|
||||
# value = self.value_net(latent_vf)
|
||||
return log_prob, distribution.entropy()
|
||||
|
||||
def reset_noise(self):
|
||||
"""
|
||||
Sample new weights for the exploration matrix.
|
||||
Sample new weights for the exploration matrix, when using SDE.
|
||||
"""
|
||||
self.action_dist.sample_weights(self.log_std)
|
||||
|
||||
def forward(self, obs, deterministic=True):
|
||||
if self.use_sde:
|
||||
latent_pi = self.latent_pi(obs)
|
||||
latent_pi, latent_sde = self._get_latent(obs)
|
||||
if deterministic:
|
||||
return self.actor_net(latent_pi)
|
||||
noise = self.action_dist.get_noise(latent_pi)
|
||||
return self.mu(latent_pi)
|
||||
|
||||
noise = self.action_dist.get_noise(latent_sde)
|
||||
if self.clip_noise is not None:
|
||||
noise = th.clamp(noise, -self.clip_noise, self.clip_noise)
|
||||
# TODO: Replace with squashing -> need to account for that in the sde update
|
||||
# -> set squash_out=True in the action_dist?
|
||||
# NOTE: the clipping is done in the rollout for now
|
||||
return self.actor_net(latent_pi) + noise
|
||||
return self.mu(latent_pi) + noise
|
||||
# action, _ = self._get_action_dist_from_latent(latent_pi)
|
||||
# return action
|
||||
else:
|
||||
return self.actor_net(obs)
|
||||
return self.mu(obs)
|
||||
|
||||
|
||||
class Critic(BaseNetwork):
|
||||
|
|
@ -136,23 +161,50 @@ class Critic(BaseNetwork):
|
|||
return self.q_networks[0](th.cat([obs, action], dim=1))
|
||||
|
||||
|
||||
class ValueFunction(BaseNetwork):
|
||||
"""
|
||||
Value function for TD3 when doing on-policy exploration with SDE.
|
||||
|
||||
:param obs_dim: (int) Dimension of the observation
|
||||
:param net_arch: ([int]) Network architecture
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
"""
|
||||
def __init__(self, obs_dim, net_arch=None, activation_fn=nn.Tanh):
|
||||
super(ValueFunction, self).__init__()
|
||||
|
||||
if net_arch is None:
|
||||
net_arch = [64, 64]
|
||||
|
||||
vf_net = create_mlp(obs_dim, 1, net_arch, activation_fn)
|
||||
self.vf_net = nn.Sequential(*vf_net)
|
||||
|
||||
def forward(self, obs):
|
||||
return self.vf_net(obs)
|
||||
|
||||
|
||||
class TD3Policy(BasePolicy):
|
||||
"""
|
||||
Policy class (with both actor and critic) for TD3.
|
||||
|
||||
:param observation_space: (gym.spaces.Space) Observation space
|
||||
:param action_dim: (gym.spaces.Space) Action space
|
||||
:param action_space: (gym.spaces.Space) Action space
|
||||
:param learning_rate: (callable) Learning rate schedule (could be constant)
|
||||
:param net_arch: ([int or dict]) The specification of the policy and value networks.
|
||||
:param device: (str or th.device) Device on which the code should run.
|
||||
:param activation_fn: (nn.Module) Activation function
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration or not
|
||||
:param log_std_init: (float) Initial value for the log standard deviation
|
||||
:param sde_net_arch: ([int]) Network architecture for extracting features
|
||||
when using SDE. If None, the latent features from the policy will be used.
|
||||
Pass an empty list to use the states as features.
|
||||
:param use_expln: (bool) Use `expln()` function instead of `exp()` when using SDE to ensure
|
||||
a positive standard deviation (cf paper). It allows to keep variance
|
||||
above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
|
||||
"""
|
||||
def __init__(self, observation_space, action_space,
|
||||
learning_rate, net_arch=None, device='cpu',
|
||||
activation_fn=nn.ReLU, use_sde=False, log_std_init=-2,
|
||||
clip_noise=None, lr_sde=3e-4):
|
||||
activation_fn=nn.ReLU, use_sde=False, log_std_init=-3,
|
||||
clip_noise=None, lr_sde=3e-4, sde_net_arch=None, use_expln=False):
|
||||
super(TD3Policy, self).__init__(observation_space, action_space, device)
|
||||
|
||||
# Default network architecture, from the original paper
|
||||
|
|
@ -170,14 +222,21 @@ class TD3Policy(BasePolicy):
|
|||
'activation_fn': self.activation_fn
|
||||
}
|
||||
self.actor_kwargs = self.net_args.copy()
|
||||
self.actor_kwargs['use_sde'] = use_sde
|
||||
self.actor_kwargs['log_std_init'] = log_std_init
|
||||
self.actor_kwargs['clip_noise'] = clip_noise
|
||||
self.actor_kwargs['lr_sde'] = lr_sde
|
||||
sde_kwargs = {
|
||||
'use_sde': use_sde,
|
||||
'log_std_init': log_std_init,
|
||||
'clip_noise': clip_noise,
|
||||
'lr_sde': lr_sde,
|
||||
'sde_net_arch': sde_net_arch,
|
||||
'use_expln': use_expln
|
||||
}
|
||||
self.actor_kwargs.update(sde_kwargs)
|
||||
|
||||
self.actor, self.actor_target = None, None
|
||||
self.critic, self.critic_target = None, None
|
||||
# For SDE only
|
||||
self.use_sde = use_sde
|
||||
self.vf_net = None
|
||||
self.log_std_init = log_std_init
|
||||
self._build(learning_rate)
|
||||
|
||||
|
|
@ -192,6 +251,10 @@ class TD3Policy(BasePolicy):
|
|||
self.critic_target.load_state_dict(self.critic.state_dict())
|
||||
self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
|
||||
|
||||
if self.use_sde:
|
||||
self.vf_net = ValueFunction(self.obs_dim)
|
||||
self.actor.sde_optimizer.add_param_group({'params': self.vf_net.parameters()})
|
||||
|
||||
def reset_noise(self):
|
||||
return self.actor.reset_noise()
|
||||
|
||||
|
|
|
|||
|
|
@ -1,14 +1,10 @@
|
|||
import time
|
||||
|
||||
import torch as th
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.base_class import BaseRLModel
|
||||
from torchy_baselines.common.buffers import ReplayBuffer
|
||||
from torchy_baselines.common.evaluation import evaluate_policy
|
||||
from torchy_baselines.td3.policies import TD3Policy
|
||||
from torchy_baselines.common.vec_env import sync_envs_normalization
|
||||
|
||||
|
||||
class TD3(BaseRLModel):
|
||||
|
|
@ -40,6 +36,8 @@ class TD3(BaseRLModel):
|
|||
:param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise.
|
||||
:param use_sde: (bool) Whether to use State Dependent Exploration (SDE)
|
||||
instead of action noise exploration (default: False)
|
||||
:param sde_sample_freq: (int) Sample a new noise matrix every n steps when using SDE
|
||||
Default: -1 (only sample at the beginning of the rollout)
|
||||
:param sde_max_grad_norm: (float)
|
||||
:param sde_ent_coef: (float)
|
||||
:param sde_log_std_scheduler: (callable)
|
||||
|
|
@ -57,12 +55,13 @@ class TD3(BaseRLModel):
|
|||
policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100,
|
||||
train_freq=-1, gradient_steps=-1, n_episodes_rollout=1,
|
||||
tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5,
|
||||
use_sde=False, sde_max_grad_norm=1, sde_ent_coef=0.0, sde_log_std_scheduler=None,
|
||||
use_sde=False, sde_sample_freq=-1, sde_max_grad_norm=1, sde_ent_coef=0.0, sde_log_std_scheduler=None,
|
||||
tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0,
|
||||
seed=0, device='auto', _init_setup_model=True):
|
||||
|
||||
super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose, device,
|
||||
create_eval_env=create_eval_env, seed=seed)
|
||||
create_eval_env=create_eval_env, seed=seed,
|
||||
use_sde=use_sde, sde_sample_freq=sde_sample_freq)
|
||||
|
||||
self.buffer_size = buffer_size
|
||||
self.learning_rate = learning_rate
|
||||
|
|
@ -79,10 +78,11 @@ class TD3(BaseRLModel):
|
|||
self.target_policy_noise = target_policy_noise
|
||||
|
||||
# State Dependent Exploration
|
||||
self.use_sde = use_sde
|
||||
self.sde_max_grad_norm = sde_max_grad_norm
|
||||
self.sde_ent_coef = sde_ent_coef
|
||||
self.sde_log_std_scheduler = sde_log_std_scheduler
|
||||
self.on_policy_exploration = True
|
||||
self.sde_vf = None
|
||||
|
||||
if _init_setup_model:
|
||||
self._setup_model()
|
||||
|
|
@ -93,7 +93,8 @@ class TD3(BaseRLModel):
|
|||
self.set_random_seed(self.seed)
|
||||
self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
|
||||
self.policy = self.policy_class(self.observation_space, self.action_space,
|
||||
self.learning_rate, use_sde=self.use_sde, device=self.device, **self.policy_kwargs)
|
||||
self.learning_rate, use_sde=self.use_sde,
|
||||
device=self.device, **self.policy_kwargs)
|
||||
self.policy = self.policy.to(self.device)
|
||||
self._create_aliases()
|
||||
|
||||
|
|
@ -102,6 +103,7 @@ class TD3(BaseRLModel):
|
|||
self.actor_target = self.policy.actor_target
|
||||
self.critic = self.policy.critic
|
||||
self.critic_target = self.policy.critic_target
|
||||
self.vf_net = self.policy.vf_net
|
||||
|
||||
def select_action(self, observation, deterministic=True):
|
||||
# Normally not needed
|
||||
|
|
@ -207,25 +209,27 @@ class TD3(BaseRLModel):
|
|||
# self._update_learning_rate(self.policy.optimizer)
|
||||
|
||||
# Unpack
|
||||
obs, action, returns = [self.rollout_data[key] for key in ['observations', 'actions', 'returns']]
|
||||
obs, action, advantage, returns = [self.rollout_data[key] for key in
|
||||
['observations', 'actions', 'advantage', 'returns']]
|
||||
|
||||
# TODO: avoid second computation of everything because of the gradient
|
||||
log_prob, entropy = self.actor.evaluate_actions(obs, action)
|
||||
values = self.vf_net(obs).flatten()
|
||||
|
||||
# Normalize returns
|
||||
# returns = (returns - returns.mean()) / (returns.std() + 1e-8)
|
||||
# returns = (returns - returns.mean())
|
||||
with th.no_grad():
|
||||
current_q1, current_q2 = self.critic(obs, action)
|
||||
# Alternatively use the q value
|
||||
returns = (returns - th.min(current_q1, current_q2))
|
||||
# Normalize advantage
|
||||
# if self.normalize_advantage:
|
||||
# advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
|
||||
|
||||
policy_loss = -(returns * log_prob).mean()
|
||||
# Value loss using the TD(gae_lambda) target
|
||||
value_loss = F.mse_loss(returns, values)
|
||||
|
||||
# A2C loss
|
||||
policy_loss = -(advantage * log_prob).mean()
|
||||
|
||||
# Entropy loss favor exploration
|
||||
entropy_loss = -th.mean(entropy)
|
||||
|
||||
loss = policy_loss + self.sde_ent_coef * entropy_loss
|
||||
vf_coef = 0.5
|
||||
loss = policy_loss + self.sde_ent_coef * entropy_loss + vf_coef * value_loss
|
||||
|
||||
# Optimization step
|
||||
self.actor.sde_optimizer.zero_grad()
|
||||
|
|
@ -284,15 +288,9 @@ class TD3(BaseRLModel):
|
|||
gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps
|
||||
self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay)
|
||||
|
||||
# Evaluate episode
|
||||
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
|
||||
timesteps_since_eval %= eval_freq
|
||||
sync_envs_normalization(self.env, eval_env)
|
||||
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
|
||||
evaluations.append(mean_reward)
|
||||
if self.verbose > 0:
|
||||
print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))
|
||||
print("FPS: {:.2f}".format(self.num_timesteps / (time.time() - self.start_time)))
|
||||
# Evaluate the agent
|
||||
timesteps_since_eval = self._eval_policy(eval_freq, eval_env, n_eval_episodes,
|
||||
timesteps_since_eval, deterministic=True)
|
||||
|
||||
return self
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue