saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-17 22:30:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #5 from Antonin-Raffin/feat/td3-sde

Browse source 
Off-Policy State Dependent Exploration
This commit is contained in:
Raffin, Antonin 2019-12-17 11:09:01 +01:00 committed by GitHub Enterprise
commit 8874b9dd6b
19 changed files with 820 additions and 242 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
README.md
View file

@ -21,6 +21,10 @@ TODO:
- save/load - save/load
- better predict - better predict
- complete logger - complete logger
- SDE: learn the feature extractor?
- Refactor: buffer with numpy array instead of pytorch
- Refactor: remove duplicated code for evaluation
- plotting? -> zoo
Later: Later:
- get_parameters / set_parameters - get_parameters / set_parameters

2
setup.py
View file

@ -34,7 +34,7 @@ setup(name='torchy_baselines',
license="MIT", license="MIT",
long_description="", long_description="",
long_description_content_type='text/markdown', long_description_content_type='text/markdown',
version="0.0.5a", version="0.0.6a",
) )
# python setup.py sdist # python setup.py sdist

19
tests/test_sde.py
View file

@ -4,7 +4,7 @@ import gym
import torch as th import torch as th
from torch.distributions import Normal from torch.distributions import Normal
from torchy_baselines import A2C from torchy_baselines import A2C, TD3
from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize from torchy_baselines.common.vec_env import DummyVecEnv, VecNormalize
from torchy_baselines.common.monitor import Monitor from torchy_baselines.common.monitor import Monitor
@ -58,3 +58,20 @@ def test_state_dependent_noise(model_class):
model = model_class('MlpPolicy', env, n_steps=200, use_sde=True, ent_coef=0.00, verbose=1, learning_rate=3e-4, 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) 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) 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):
model = model_class('MlpPolicy', 'Pendulum-v0', use_sde=True, seed=None, create_eval_env=True,
verbose=1, policy_kwargs=dict(log_std_init=-2))
model.learn(total_timesteps=int(1000), eval_freq=500)
def test_scheduler():
def scheduler(progress):
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)
model.learn(total_timesteps=int(1000), eval_freq=500)
assert th.isclose(model.actor.log_std, th.ones_like(model.actor.log_std)).all()

14
tests/test_vec_normalize.py
View file

@ -1,9 +1,11 @@
import gym import gym
import pytest
import numpy as np import numpy as np
from torchy_baselines.common.running_mean_std import RunningMeanStd 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.dummy_vec_env import DummyVecEnv
from torchy_baselines.common.vec_env.vec_normalize import VecNormalize from torchy_baselines.common.vec_env.vec_normalize import VecNormalize
from torchy_baselines import CEMRL, SAC, TD3
ENV_ID = 'Pendulum-v0' ENV_ID = 'Pendulum-v0'
@ -39,3 +41,15 @@ def test_vec_env():
actions = [env.action_space.sample()] actions = [env.action_space.sample()]
obs, _, done, _ = env.step(actions) obs, _, done, _ = env.step(actions)
assert np.max(obs) <= 10 assert np.max(obs) <= 10
@pytest.mark.parametrize("model_class", [SAC, TD3, CEMRL])
def test_offpolicy_normalization(model_class):
env = DummyVecEnv([lambda: gym.make(ENV_ID)])
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.)
model = model_class('MlpPolicy', env, verbose=1)
model.learn(total_timesteps=1000, eval_env=eval_env, eval_freq=500)

2
torchy_baselines/__init__.py
View file

@ -4,4 +4,4 @@ from torchy_baselines.ppo import PPO
from torchy_baselines.sac import SAC from torchy_baselines.sac import SAC
from torchy_baselines.td3 import TD3 from torchy_baselines.td3 import TD3
__version__ = "0.0.5a" __version__ = "0.0.6a"

3
torchy_baselines/a2c/a2c.py
View file

@ -25,6 +25,7 @@ class A2C(PPO):
(i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel) (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel)
:param gamma: (float) Discount factor :param gamma: (float) Discount factor
:param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator :param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
Equivalent to classic advantage when set to 1.
:param ent_coef: (float) Entropy coefficient for the loss calculation :param ent_coef: (float) Entropy coefficient for the loss calculation
:param vf_coef: (float) Value function coefficient for the loss calculation :param vf_coef: (float) Value function coefficient for the loss calculation
:param max_grad_norm: (float) The maximum value for the gradient clipping :param max_grad_norm: (float) The maximum value for the gradient clipping
@ -92,7 +93,7 @@ class A2C(PPO):
action = action.long().flatten() action = action.long().flatten()
# TODO: avoid second computation of everything because of the gradient # TODO: avoid second computation of everything because of the gradient
values, log_prob, entropy = self.policy.get_policy_stats(obs, action) values, log_prob, entropy = self.policy.evaluate_actions(obs, action)
values = values.flatten() values = values.flatten()
# Normalize advantage (not present in the original implementation) # Normalize advantage (not present in the original implementation)

80
torchy_baselines/cem_rl/cem.py
View file

@ -5,33 +5,42 @@ import numpy as np
# or https://github.com/facebookresearch/nevergrad # or https://github.com/facebookresearch/nevergrad
class CEM(object): class CEM(object):
"""
Cross-entropy method with diagonal covariance (separable CEM)
""" """
Cross-entropy method with diagonal covariance (separable CEM).
def __init__(self, num_params, :param num_params: (int) Number of parameters per individual (dimension of the problem)
mu_init=None, :param mu_init: (np.ndarray) Initial mean of the population distribution
sigma_init=1e-3, Taken to be zero if None is passed.
pop_size=256, :param sigma_init: (float) Initial standard deviation of the population distribution
damp=1e-3, :param pop_size: (int) Number of individuals in the population
damp_limit=1e-5, :param damping_init: (float) Initial value of damping for preventing from early convergence.
parents=None, :param damping_final: (float) Final value of damping
elitism=False, :param parents: (int) Number of parents used to compute the new distribution
antithetic=False): of individuals.
:param elitism: (bool) Keep the best known individual in the population
:param antithetic: (bool) Use a finite difference like method for sampling
(mu + epsilon, mu - epsilon)
"""
def __init__(self, num_params, mu_init=None, sigma_init=1e-3,
pop_size=256, damping_init=1e-3, damping_final=1e-5,
parents=None, elitism=False, antithetic=False):
super(CEM, self).__init__() super(CEM, self).__init__()
# misc
self.num_params = num_params self.num_params = num_params
# distribution parameters # Distribution parameters
if mu_init is None: if mu_init is None:
self.mu = np.zeros(self.num_params) self.mu = np.zeros(self.num_params)
else: else:
self.mu = np.array(mu_init) self.mu = np.array(mu_init)
self.sigma = sigma_init self.sigma = sigma_init
self.damp = damp # Damping parameters
self.damp_limit = damp_limit self.damping = damping_init
self.damping_final = damping_final
# Exponential moving average decay for damping
self.tau = 0.95 self.tau = 0.95
# Covariance matrix, here only the diagonal
self.cov = self.sigma * np.ones(self.num_params) self.cov = self.sigma * np.ones(self.num_params)
# elite stuff # elite stuff
@ -39,16 +48,20 @@ class CEM(object):
self.elite = np.sqrt(self.sigma) * np.random.rand(self.num_params) self.elite = np.sqrt(self.sigma) * np.random.rand(self.num_params)
self.elite_score = None self.elite_score = None
# sampling stuff # sampling parameters
self.pop_size = pop_size self.pop_size = pop_size
self.antithetic = antithetic self.antithetic = antithetic
if self.antithetic: if self.antithetic:
assert (self.pop_size % 2 == 0), "Population size must be even" assert (self.pop_size % 2 == 0), "Population size must be even"
if parents is None or parents <= 0: if parents is None or parents <= 0:
self.parents = pop_size // 2 self.parents = pop_size // 2
else: else:
self.parents = parents self.parents = parents
# Weighting for computing the new mean of the distributions
# from the parents. The better the individual, the higher the weight
self.weights = np.array([np.log((self.parents + 1) / i) self.weights = np.array([np.log((self.parents + 1) / i)
for i in range(1, self.parents + 1)]) for i in range(1, self.parents + 1)])
self.weights /= self.weights.sum() self.weights /= self.weights.sum()
@ -56,43 +69,56 @@ class CEM(object):
def ask(self, pop_size): def ask(self, pop_size):
""" """
Returns a list of candidates parameters Returns a list of candidates parameters
:param pop_size: (int)
:return: ([np.ndarray])
""" """
if self.antithetic and not pop_size % 2: if self.antithetic and not pop_size % 2:
epsilon_half = np.random.randn(pop_size // 2, self.num_params) epsilon_half = np.random.randn(pop_size // 2, self.num_params)
epsilon = np.concatenate([epsilon_half, - epsilon_half]) epsilon = np.concatenate([epsilon_half, - epsilon_half])
else: else:
epsilon = np.random.randn(pop_size, self.num_params) epsilon = np.random.randn(pop_size, self.num_params)
inds = self.mu + epsilon * np.sqrt(self.cov) individuals = self.mu + epsilon * np.sqrt(self.cov)
if self.elitism:
inds[-1] = self.elite
return inds # Keep the best known individual in the population
if self.elitism:
individuals[-1] = self.elite
return individuals
def tell(self, solutions, scores): def tell(self, solutions, scores):
""" """
Updates the distribution Updates the distribution
:param solutions: ([np.ndarray])
:param scores: ([float]) episode reward.
""" """
# Convert rewards (we want to maximize) to cost (we want to minimize)
scores = np.array(scores) scores = np.array(scores)
scores *= -1 scores *= -1
# Sort the individuals by fitness
idx_sorted = np.argsort(scores) idx_sorted = np.argsort(scores)
old_mu = self.mu old_mu = self.mu
self.damp = self.damp * self.tau + (1 - self.tau) * self.damp_limit # Update damping using a moving average
self.damping = self.damping * self.tau + (1 - self.tau) * self.damping_final
# self.mu = self.weights @ solutions[idx_sorted[:self.parents]] # self.mu = self.weights @ solutions[idx_sorted[:self.parents]]
self.mu = self.weights.dot(solutions[idx_sorted[:self.parents]]) self.mu = self.weights.dot(solutions[idx_sorted[:self.parents]])
# CMA-ES style would be to use the new mean here
z = (solutions[idx_sorted[:self.parents]] - old_mu) z = (solutions[idx_sorted[:self.parents]] - old_mu)
self.cov = 1 / self.parents * self.weights.dot(z * z) + self.damp * np.ones(self.num_params) self.cov = 1 / self.parents * self.weights.dot(z * z) + self.damping * np.ones(self.num_params)
# Retrieve the best individual
self.elite = solutions[idx_sorted[0]] self.elite = solutions[idx_sorted[0]]
self.elite_score = scores[idx_sorted[0]] self.elite_score = scores[idx_sorted[0]]
# print(self.cov)
def get_distrib_params(self): def get_distrib_params(self):
""" """
Returns the parameters of the distrubtion: Returns the parameters of the distribution:
the mean and sigma the mean and standard deviation.
:return: (np.ndarray, np.ndarray)
""" """
return np.copy(self.mu), np.copy(self.cov) return np.copy(self.mu), np.copy(self.cov)

66
torchy_baselines/cem_rl/cem_rl.py
View file

@ -5,21 +5,53 @@ import torch as th
from torchy_baselines.cem_rl.cem import CEM from torchy_baselines.cem_rl.cem import CEM
from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.common.evaluation import evaluate_policy
from torchy_baselines.td3.td3 import TD3 from torchy_baselines.td3.td3 import TD3
from torchy_baselines.common.vec_env import sync_envs_normalization
class CEMRL(TD3): class CEMRL(TD3):
""" """
Implementation of CEM-RL Implementation of CEM-RL, in fact CEM combined with TD3.
Paper: https://arxiv.org/abs/1810.01222 Paper: https://arxiv.org/abs/1810.01222
Code: https://github.com/apourchot/CEM-RL Code: https://github.com/apourchot/CEM-RL
"""
:param policy: (TD3Policy or str) The policy model to use (MlpPolicy, CnnPolicy, ...)
:param env: (Gym environment or str) The environment to learn from (if registered in Gym, can be str)
:param sigma_init: (float) Initial standard deviation of the population distribution
:param pop_size: (int) Number of individuals in the population
:param damping_init: (float) Initial value of damping for preventing from early convergence.
:param damping_final: (float) Final value of damping
:param elitism: (bool) Keep the best known individual in the population
:param n_grad: (int) Number of individuals that will receive a gradient update.
Half of the population size in the paper.
:param buffer_size: (int) size of the replay buffer
:param learning_rate: (float or callable) learning rate for adam optimizer,
the same learning rate will be used for all networks (Q-Values and Actor networks)
it can be a function of the current progress (from 1 to 0)
:param policy_delay: (int) Policy and target networks will only be updated once every policy_delay steps
per training steps. The Q values will be updated policy_delay more often (update every training step).
:param learning_starts: (int) how many steps of the model to collect transitions for before learning starts
:param gamma: (float) the discount factor
:param batch_size: (int) Minibatch size for each gradient update
:param tau: (float) the soft update coefficient ("polyak update" of the target networks, between 0 and 1)
:param action_noise: (ActionNoise) the action noise type. Cf common.noise for the different action noise type.
:param target_policy_noise: (float) Standard deviation of gaussian noise added to target policy
(smoothing noise)
:param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise.
: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
:param verbose: (int) the verbosity level: 0 none, 1 training information, 2 tensorflow debug
:param seed: (int) Seed for the pseudo random generators
:param device: (str or th.device) Device (cpu, cuda, ...) on which the code should be run.
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, sigma_init=1e-3, pop_size=10, def __init__(self, policy, env, sigma_init=1e-3, pop_size=10,
damp=1e-3, damp_limit=1e-5, elitism=False, n_grad=5, damping_init=1e-3, damping_final=1e-5, elitism=False, n_grad=5,
policy_delay=2, batch_size=100, buffer_size=int(1e6), learning_rate=1e-3, policy_delay=2,
buffer_size=int(1e6), learning_rate=1e-3, learning_starts=100, gamma=0.99, batch_size=100, tau=0.005,
action_noise=None, learning_starts=100, tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5,
n_episodes_rollout=1, update_style='original', n_episodes_rollout=1, update_style='original',
tensorboard_log=None, create_eval_env=False, tensorboard_log=None, create_eval_env=False,
policy_kwargs=None, verbose=0, seed=0, device='auto', policy_kwargs=None, verbose=0, seed=0, device='auto',
@ -27,18 +59,21 @@ class CEMRL(TD3):
super(CEMRL, self).__init__(policy, env, super(CEMRL, self).__init__(policy, env,
buffer_size=buffer_size, learning_rate=learning_rate, seed=seed, device=device, buffer_size=buffer_size, learning_rate=learning_rate, seed=seed, device=device,
action_noise=action_noise, learning_starts=learning_starts, action_noise=action_noise, target_policy_noise=target_policy_noise,
n_episodes_rollout=n_episodes_rollout, tau=tau, target_noise_clip=target_noise_clip, learning_starts=learning_starts,
n_episodes_rollout=n_episodes_rollout, tau=tau, gamma=gamma,
policy_kwargs=policy_kwargs, verbose=verbose, policy_kwargs=policy_kwargs, verbose=verbose,
policy_delay=policy_delay, batch_size=batch_size, policy_delay=policy_delay, batch_size=batch_size,
create_eval_env=create_eval_env, create_eval_env=create_eval_env, tensorboard_log=tensorboard_log,
_init_setup_model=False) _init_setup_model=False)
# Evolution strategy method that follows cma-es interface (ask-tell)
# for now, only CEM is implemented
self.es = None self.es = None
self.sigma_init = sigma_init self.sigma_init = sigma_init
self.pop_size = pop_size self.pop_size = pop_size
self.damp = damp self.damping_init = damping_init
self.damp_limit = damp_limit self.damping_final = damping_final
self.elitism = elitism self.elitism = elitism
self.n_grad = n_grad self.n_grad = n_grad
self.es_params = None self.es_params = None
@ -52,7 +87,7 @@ class CEMRL(TD3):
super(CEMRL, self)._setup_model() super(CEMRL, self)._setup_model()
params_vector = self.actor.parameters_to_vector() params_vector = self.actor.parameters_to_vector()
self.es = CEM(len(params_vector), mu_init=params_vector, self.es = CEM(len(params_vector), mu_init=params_vector,
sigma_init=self.sigma_init, damp=self.damp, damp_limit=self.damp_limit, sigma_init=self.sigma_init, damping_init=self.damping_init, damping_final=self.damping_final,
pop_size=self.pop_size, antithetic=not self.pop_size % 2, parents=self.pop_size // 2, pop_size=self.pop_size, antithetic=not self.pop_size % 2, parents=self.pop_size // 2,
elitism=self.elitism) elitism=self.elitism)
@ -103,7 +138,7 @@ class CEMRL(TD3):
n_training_steps = 2 * (actor_steps // self.n_grad) n_training_steps = 2 * (actor_steps // self.n_grad)
for it in range(n_training_steps): for it in range(n_training_steps):
# Sample replay buffer # Sample replay buffer
replay_data = self.replay_buffer.sample(self.batch_size) replay_data = self.replay_buffer.sample(self.batch_size, env=self._vec_normalize_env)
self.train_critic(replay_data=replay_data) self.train_critic(replay_data=replay_data)
# Delayed policy updates # Delayed policy updates
@ -118,6 +153,7 @@ class CEMRL(TD3):
timesteps_since_eval %= eval_freq timesteps_since_eval %= eval_freq
self.actor.load_from_vector(self.es.mu) 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, _ = evaluate_policy(self, eval_env, n_eval_episodes)
evaluations.append(mean_reward) evaluations.append(mean_reward)
@ -150,10 +186,6 @@ class CEMRL(TD3):
actor_steps += episode_timesteps actor_steps += episode_timesteps
self.fitnesses.append(episode_reward) self.fitnesses.append(episode_reward)
if self.verbose > 1:
print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
self.num_timesteps, episode_num, episode_timesteps, episode_reward))
self._update_current_progress(self.num_timesteps, total_timesteps) self._update_current_progress(self.num_timesteps, total_timesteps)
self.es.tell(self.es_params, self.fitnesses) self.es.tell(self.es_params, self.fitnesses)
timesteps_since_eval += actor_steps timesteps_since_eval += actor_steps

150
torchy_baselines/common/base_class.py
View file

@ -11,7 +11,7 @@ import numpy as np
from torchy_baselines.common.policies import get_policy_from_name 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.utils import set_random_seed, get_schedule_fn, update_learning_rate
from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize
from torchy_baselines.common.monitor import Monitor from torchy_baselines.common.monitor import Monitor
from torchy_baselines.common import logger from torchy_baselines.common import logger
from torchy_baselines.common.save_util import data_to_json, json_to_data from torchy_baselines.common.save_util import data_to_json, json_to_data
@ -24,13 +24,19 @@ class BaseRLModel(object):
:param policy: (BasePolicy) Policy object :param policy: (BasePolicy) Policy object
:param env: (Gym environment) The environment to learn from :param env: (Gym environment) The environment to learn from
(if registered in Gym, can be str. Can be None for loading trained models) (if registered in Gym, can be str. Can be None for loading trained models)
:param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug
:param policy_base: (BasePolicy) the base policy used by this method :param policy_base: (BasePolicy) the base policy used by this method
:param device: (str or th.device) Device on which the code should. :param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
:param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug
:param device: (str or th.device) Device on which the code should run.
By default, it will try to use a Cuda compatible device and fallback to cpu By default, it will try to use a Cuda compatible device and fallback to cpu
if it is not possible. if it is not possible.
:param support_multi_env: (bool) Whether the algorithm supports training
with multiple environments (as in A2C)
: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 monitor_wrapper: (bool) When creating an environment, whether to wrap it :param monitor_wrapper: (bool) When creating an environment, whether to wrap it
or not in a Monitor wrapper. or not in a Monitor wrapper.
:param seed: (int) Seed for the pseudo random generators
""" """
__metaclass__ = ABCMeta __metaclass__ = ABCMeta
@ -50,6 +56,8 @@ class BaseRLModel(object):
print("Using {} device".format(self.device)) print("Using {} device".format(self.device))
self.env = env self.env = env
# get VecNormalize object if needed
self._vec_normalize_env = unwrap_vec_normalize(env)
self.verbose = verbose self.verbose = verbose
self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs
self.observation_space = None self.observation_space = None
@ -60,10 +68,14 @@ class BaseRLModel(object):
self.replay_buffer = None self.replay_buffer = None
self.seed = seed self.seed = seed
self.action_noise = None self.action_noise = None
# Used for SDE only
self.rollout_data = None
self.use_sde = False
# Track the training progress (from 1 to 0) # Track the training progress (from 1 to 0)
# this is used to update the learning rate # this is used to update the learning rate
self._current_progress = 1 self._current_progress = 1
# Create and wrap the env if needed
if env is not None: if env is not None:
if isinstance(env, str): if isinstance(env, str):
if create_eval_env: if create_eval_env:
@ -93,6 +105,12 @@ class BaseRLModel(object):
" environment.") " environment.")
def _get_eval_env(self, eval_env): def _get_eval_env(self, eval_env):
"""
Return the environment that will be used for evaluation.
:param eval_env: (gym.Env or VecEnv)
:return: (VecEnv)
"""
if eval_env is None: if eval_env is None:
eval_env = self.eval_env eval_env = self.eval_env
@ -106,6 +124,9 @@ class BaseRLModel(object):
""" """
Rescale the action from [low, high] to [-1, 1] Rescale the action from [low, high] to [-1, 1]
(no need for symmetric action space) (no need for symmetric action space)
:param action: (np.ndarray)
:return: (np.ndarray)
""" """
low, high = self.action_space.low, self.action_space.high low, high = self.action_space.low, self.action_space.high
return 2.0 * ((action - low) / (high - low)) - 1.0 return 2.0 * ((action - low) / (high - low)) - 1.0
@ -114,6 +135,9 @@ class BaseRLModel(object):
""" """
Rescale the action from [-1, 1] to [low, high] Rescale the action from [-1, 1] to [low, high]
(no need for symmetric action space) (no need for symmetric action space)
:param scaled_action: (np.ndarray)
:return: (np.ndarray)
""" """
low, high = self.action_space.low, self.action_space.high low, high = self.action_space.low, self.action_space.high
return low + (0.5 * (scaled_action + 1.0) * (high - low)) return low + (0.5 * (scaled_action + 1.0) * (high - low))
@ -126,7 +150,7 @@ class BaseRLModel(object):
""" """
Compute current progress (from 1 to 0) Compute current progress (from 1 to 0)
:param num_timesteps: (int) :param num_timesteps: (int) current number of timesteps
:param total_timesteps: (int) :param total_timesteps: (int)
""" """
self._current_progress = 1.0 - float(num_timesteps) / float(total_timesteps) self._current_progress = 1.0 - float(num_timesteps) / float(total_timesteps)
@ -162,7 +186,7 @@ class BaseRLModel(object):
""" """
returns the current environment (can be None if not defined) returns the current environment (can be None if not defined)
:return: (Gym Environment) The current environment :return: (gym.Env) The current environment
""" """
return self.env return self.env
@ -190,7 +214,7 @@ class BaseRLModel(object):
- observation_space - observation_space
- action_space - action_space
:param env: (Gym Environment) The environment for learning a policy :param env: (gym.Env) The environment for learning a policy
""" """
if self.check_env(env, self.observation_space, self.action_space) is False: 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")
@ -251,12 +275,14 @@ class BaseRLModel(object):
Return a trained model. Return a trained model.
:param total_timesteps: (int) The total number of samples to train on :param total_timesteps: (int) The total number of samples to train on
:param seed: (int) The initial seed for training, if None: keep current seed
:param callback: (function (dict, dict)) -> boolean function called at every steps with state of the algorithm. :param callback: (function (dict, dict)) -> boolean function called at every steps with state of the algorithm.
It takes the local and global variables. If it returns False, training is aborted. It takes the local and global variables. If it returns False, training is aborted.
:param log_interval: (int) The number of timesteps before logging. :param log_interval: (int) The number of timesteps before logging.
:param tb_log_name: (str) the name of the run for tensorboard log :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 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)
:return: (BaseRLModel) the trained model :return: (BaseRLModel) the trained model
""" """
pass pass
@ -407,21 +433,33 @@ class BaseRLModel(object):
self.eval_env.seed(seed) self.eval_env.seed(seed)
def _setup_learn(self, eval_env): def _setup_learn(self, eval_env):
"""
Initialize different variables needed for training.
:param eval_env: (gym.Env or VecEnv)
:return: (int, int, [float], np.ndarray, VecEnv)
"""
self.start_time = time.time() self.start_time = time.time()
self.ep_info_buffer = deque(maxlen=100) self.ep_info_buffer = deque(maxlen=100)
if self.action_noise is not None: if self.action_noise is not None:
self.action_noise.reset() self.action_noise.reset()
timesteps_since_eval, episode_num = 0, 0 timesteps_since_eval, episode_num = 0, 0
evaluations = [] evaluations = []
if eval_env is not None and self.seed is not None: if eval_env is not None and self.seed is not None:
eval_env.seed(self.seed) eval_env.seed(self.seed)
eval_env = self._get_eval_env(eval_env) eval_env = self._get_eval_env(eval_env)
obs = self.env.reset() obs = self.env.reset()
return timesteps_since_eval, episode_num, evaluations, obs, eval_env return timesteps_since_eval, episode_num, evaluations, obs, eval_env
def _update_info_buffer(self, infos): def _update_info_buffer(self, infos):
""" """
Retrieve reward and episode length if using Monitor wrapper. Retrieve reward and episode length and update the buffer
if using Monitor wrapper.
:param infos: ([dict]) :param infos: ([dict])
""" """
for info in infos: for info in infos:
@ -434,13 +472,39 @@ class BaseRLModel(object):
learning_starts=0, num_timesteps=0, learning_starts=0, num_timesteps=0,
replay_buffer=None, obs=None, replay_buffer=None, obs=None,
episode_num=0, log_interval=None): episode_num=0, log_interval=None):
"""
Collect rollout using the current policy (and possibly fill the replay buffer)
TODO: move this method to off-policy base class.
:param env: (VecEnv)
:param n_episodes: (int)
:param n_steps: (int)
:param action_noise: (ActionNoise)
:param deterministic: (bool)
:param callback: (callable)
:param learning_starts: (int)
:param num_timesteps: (int)
:param replay_buffer: (ReplayBuffer)
:param obs: (np.ndarray)
:param episode_num: (int)
:param log_interval: (int)
"""
episode_rewards = [] episode_rewards = []
total_timesteps = [] total_timesteps = []
total_steps, total_episodes = 0, 0 total_steps, total_episodes = 0, 0
assert isinstance(env, VecEnv) assert isinstance(env, VecEnv)
assert env.num_envs == 1 assert env.num_envs == 1
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
obs_ = self._vec_normalize_env.get_original_obs()
self.rollout_data = None
if self.use_sde:
self.actor.reset_noise()
# Reset rollout data
self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']}
while total_steps < n_steps or total_episodes < n_episodes: while total_steps < n_steps or total_episodes < n_episodes:
done = False done = False
# Reset environment: not needed for VecEnv # Reset environment: not needed for VecEnv
@ -450,20 +514,29 @@ class BaseRLModel(object):
while not done: while not done:
# Select action randomly or according to policy # Select action randomly or according to policy
if num_timesteps < learning_starts: if num_timesteps < learning_starts:
action = np.array([self.action_space.sample()]) # Warmup phase
unscaled_action = np.array([self.action_space.sample()])
else: else:
action = self.predict(obs, deterministic=deterministic) unscaled_action = self.predict(obs, deterministic=not self.use_sde)
# Rescale the action from [low, high] to [-1, 1] # Rescale the action from [low, high] to [-1, 1]
action = self.scale_action(action) scaled_action = self.scale_action(unscaled_action)
if self.use_sde:
# When using SDE, the action can be out of bounds
# TODO: fix with squashing and account for that in the proba distribution
clipped_action = np.clip(scaled_action, -1, 1)
else:
clipped_action = scaled_action
# Add noise to the action (improve exploration) # Add noise to the action (improve exploration)
if action_noise is not None: if action_noise is not None:
# NOTE: in the original implementation of TD3, the noise was applied to the unscaled action # NOTE: in the original implementation of TD3, the noise was applied to the unscaled action
action = np.clip(action + action_noise(), -1, 1) # Update(October 2019): Not anymore
clipped_action = np.clip(clipped_action + action_noise(), -1, 1)
# Rescale and perform action # Rescale and perform action
new_obs, reward, done, infos = env.step(self.unscale_action(action)) new_obs, reward, done, infos = env.step(self.unscale_action(clipped_action))
done_bool = [float(done[0])] done_bool = [float(done[0])]
episode_reward += reward episode_reward += reward
@ -473,14 +546,34 @@ class BaseRLModel(object):
# Store data in replay buffer # Store data in replay buffer
if replay_buffer is not None: if replay_buffer is not None:
replay_buffer.add(obs, new_obs, action, reward, done_bool) # Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs()
reward_ = self._vec_normalize_env.get_original_reward()
else:
# Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, reward
replay_buffer.add(obs_, new_obs_, clipped_action, reward_, done_bool)
if self.rollout_data is not None:
# Assume only one env
self.rollout_data['observations'].append(obs[0].copy())
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 = new_obs obs = new_obs
# Save the true unnormalized observation
# otherwise obs_ = self._vec_normalize_env.unnormalize_obs(obs)
# is a good approximation
if self._vec_normalize_env is not None:
obs_ = new_obs_
num_timesteps += 1 num_timesteps += 1
episode_timesteps += 1 episode_timesteps += 1
total_steps += 1 total_steps += 1
if n_steps > 0 and total_steps >= n_steps: if 0 < n_steps <= total_steps:
break break
if done: if done:
@ -495,29 +588,46 @@ class BaseRLModel(object):
episode_num + total_episodes) % log_interval == 0: episode_num + total_episodes) % log_interval == 0:
fps = int(num_timesteps / (time.time() - self.start_time)) fps = int(num_timesteps / (time.time() - self.start_time))
logger.logkv("episodes", episode_num + total_episodes) logger.logkv("episodes", episode_num + total_episodes)
# logger.logkv("mean 100 episode reward", mean_reward)
if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0: if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer])) logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer])) logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer]))
# logger.logkv("n_updates", n_updates) # logger.logkv("n_updates", n_updates)
# logger.logkv("current_lr", current_lr)
logger.logkv("fps", fps) logger.logkv("fps", fps)
logger.logkv('time_elapsed', int(time.time() - self.start_time)) logger.logkv('time_elapsed', int(time.time() - self.start_time))
logger.logkv("total timesteps", num_timesteps) logger.logkv("total timesteps", num_timesteps)
if self.use_sde:
logger.logkv("std", (self.actor.get_std()).mean().item())
logger.dumpkvs() logger.dumpkvs()
mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0 mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0
# Post processing
if self.rollout_data is not None:
for key in ['observations', 'actions', 'rewards', 'dones']:
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
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]
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
return mean_reward, total_steps, total_episodes, obs return mean_reward, total_steps, total_episodes, obs
@staticmethod @staticmethod
def _save_to_file_zip(save_path, data=None, params=None, opt_params=None): def _save_to_file_zip(save_path, data=None, params=None, opt_params=None):
"""Save model to a zip archive """Save model to a zip archive
:param save_path: (str) Where to store the model :param save_path: (str) Where to store the model
:param data: (dict) Class parameters being stored :param data: (dict) Class parameters being stored
:param params: (dict) Model parameters being stored expected to be state_dict :param params: (dict) Model parameters being stored expected to be state_dict
:param opt_params: (dict) Optimizer parameters being stored expected to contain an entry for every :param opt_params: (dict) Optimizer parameters being stored expected to contain an entry for every
optimizer with its name and the state_dict optimizer with its name and the state_dict
""" """
# data/params can be None, so do not # data/params can be None, so do not
@ -551,7 +661,7 @@ class BaseRLModel(object):
""" """
Returns the names of the parameters that should be excluded by default Returns the names of the parameters that should be excluded by default
when saving the model. when saving the model.
:return: ([str]) List of parameters that should be excluded from save :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"]

97
torchy_baselines/common/buffers.py
View file

@ -1,8 +1,19 @@
import numpy as np import numpy as np
import torch as th import torch as th
from torchy_baselines.common.vec_env import unwrap_vec_normalize
class BaseBuffer(object): class BaseBuffer(object):
"""
Base class that represent a buffer (rollout or replay)
:param buffer_size: (int) Max number of element in the buffer
:param obs_dim: (int) Dimension of the observation
:param action_dim: (int) Dimension of the action space
:param device: (th.device)
:param n_envs: (int) Number of parallel environments
"""
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1): def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1):
super(BaseBuffer, self).__init__() super(BaseBuffer, self).__init__()
self.buffer_size = buffer_size self.buffer_size = buffer_size
@ -29,35 +40,68 @@ class BaseBuffer(object):
return tensor.transpose(0, 1).reshape(shape[0] * shape[1], *shape[2:]) return tensor.transpose(0, 1).reshape(shape[0] * shape[1], *shape[2:])
def size(self): def size(self):
"""
:return: (int) The current size of the buffer
"""
if self.full: if self.full:
return self.buffer_size return self.buffer_size
return self.pos return self.pos
def get_pos(self):
return self.pos
def add(self, *args, **kwargs): def add(self, *args, **kwargs):
"""
Add elements to the buffer.
"""
raise NotImplementedError() raise NotImplementedError()
def reset(self): def reset(self):
"""
Reset the buffer.
"""
self.pos = 0 self.pos = 0
self.full = False self.full = False
def sample(self, batch_size): def sample(self, batch_size, env=None):
"""
:param batch_size: (int) Number of element to sample
:param env: (VecNormalize) [Optional] associated gym VecEnv
to normalize the observations/rewards when sampling
"""
upper_bound = self.buffer_size if self.full else self.pos upper_bound = self.buffer_size if self.full else self.pos
batch_inds = th.LongTensor( batch_inds = th.LongTensor(
np.random.randint(0, upper_bound, size=batch_size)) np.random.randint(0, upper_bound, size=batch_size))
return self._get_samples(batch_inds) return self._get_samples(batch_inds, env=env)
def _get_samples(self, batch_inds): def _get_samples(self, batch_inds, env=None):
"""
:param batch_inds: (th.Tensor)
:param env: (gym.Env)
:return: ([th.Tensor])
"""
raise NotImplementedError() raise NotImplementedError()
def _normalize_obs(self, obs, env=None):
if env is not None:
# TODO: get rid of pytorch - numpy conversion
return th.FloatTensor(env.normalize_obs(obs.numpy()))
return obs
def _normalize_reward(self, reward, env=None):
if env is not None:
return th.FloatTensor(env.normalize_reward(reward.numpy()))
return reward
class ReplayBuffer(BaseBuffer): class ReplayBuffer(BaseBuffer):
""" """
Taken from https://github.com/apourchot/CEM-RL Replay buffer used in off-policy algorithms like SAC/TD3.
""" Adapted from from https://github.com/apourchot/CEM-RL
:param buffer_size: (int) Max number of element in the buffer
:param obs_dim: (int) Dimension of the observation
:param action_dim: (int) Dimension of the action space
:param device: (th.device)
:param n_envs: (int) Number of parallel environments
"""
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1): def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1):
super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs) super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
@ -81,15 +125,27 @@ class ReplayBuffer(BaseBuffer):
self.full = True self.full = True
self.pos = 0 self.pos = 0
def _get_samples(self, batch_inds): def _get_samples(self, batch_inds, env=None):
return (self.observations[batch_inds, 0, :].to(self.device), return (self._normalize_obs(self.observations[batch_inds, 0, :], env).to(self.device),
self.actions[batch_inds, 0, :].to(self.device), self.actions[batch_inds, 0, :].to(self.device),
self.next_observations[batch_inds, 0, :].to(self.device), self._normalize_obs(self.next_observations[batch_inds, 0, :], env).to(self.device),
self.dones[batch_inds].to(self.device), self.dones[batch_inds].to(self.device),
self.rewards[batch_inds].to(self.device)) self._normalize_reward(self.rewards[batch_inds], env).to(self.device))
class RolloutBuffer(BaseBuffer): class RolloutBuffer(BaseBuffer):
"""
Rollout buffer used in on-policy algorithms like A2C/PPO.
:param buffer_size: (int) Max number of element in the buffer
:param obs_dim: (int) Dimension of the observation
:param action_dim: (int) Dimension of the action space
:param device: (th.device)
:param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
Equivalent to classic advantage when set to 1.
:param gamma: (float) Discount factor
:param n_envs: (int) Number of parallel environments
"""
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', def __init__(self, buffer_size, obs_dim, action_dim, device='cpu',
gae_lambda=1, gamma=0.99, n_envs=1): gae_lambda=1, gamma=0.99, n_envs=1):
super(RolloutBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs) super(RolloutBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
@ -115,7 +171,10 @@ class RolloutBuffer(BaseBuffer):
def compute_returns_and_advantage(self, last_value, dones=False, use_gae=True): def compute_returns_and_advantage(self, last_value, dones=False, use_gae=True):
""" """
From Stable-Baselines PPO2 Post-processing step: compute the returns (sum of discounted rewards)
and advantage (A(s) = R - V(S)).
Adapted from Stable-Baselines PPO2.
:param last_value: (th.Tensor) :param last_value: (th.Tensor)
:param dones: ([bool]) :param dones: ([bool])
:param use_gae: (bool) Whether to use Generalized Advantage Estimation :param use_gae: (bool) Whether to use Generalized Advantage Estimation
@ -151,6 +210,16 @@ class RolloutBuffer(BaseBuffer):
self.advantages = self.returns - self.values self.advantages = self.returns - self.values
def add(self, obs, action, reward, done, value, log_prob): def add(self, obs, action, reward, done, value, log_prob):
"""
:param obs: (np.ndarray) Observation
:param action: (np.ndarray) Action
:param reward: (np.ndarray)
:param done: (np.ndarray) End of episode signal.
:param value: (th.Tensor) estimated value of the current state
following the current policy.
:param log_prob: (th.Tensor) log probability of the action
following the current policy.
"""
if len(log_prob.shape) == 0: if len(log_prob.shape) == 0:
# Reshape 0-d tensor to avoid error # Reshape 0-d tensor to avoid error
log_prob = log_prob.reshape(-1, 1) log_prob = log_prob.reshape(-1, 1)
@ -184,7 +253,7 @@ class RolloutBuffer(BaseBuffer):
yield self._get_samples(indices[start_idx:start_idx + batch_size]) yield self._get_samples(indices[start_idx:start_idx + batch_size])
start_idx += batch_size start_idx += batch_size
def _get_samples(self, batch_inds): def _get_samples(self, batch_inds, env=None):
return (self.observations[batch_inds].to(self.device), return (self.observations[batch_inds].to(self.device),
self.actions[batch_inds].to(self.device), self.actions[batch_inds].to(self.device),
self.values[batch_inds].flatten().to(self.device), self.values[batch_inds].flatten().to(self.device),

60
torchy_baselines/common/distributions.py
View file

@ -236,6 +236,8 @@ class StateDependentNoiseDistribution(Distribution):
compute the log probabilty of an action with that noise. compute the log probabilty of an action with that noise.
:param action_dim: (int) Number of continuous actions :param action_dim: (int) Number of continuous actions
:param full_std: (bool) Whether to use (n_features x n_actions) parameters
for the std instead of only (n_features,)
:param use_expln: (bool) Use `expln()` function instead of `exp()` to ensure :param use_expln: (bool) Use `expln()` function instead of `exp()` to ensure
a positive standard deviation (cf paper). It allows to keep variance 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. above zero and prevent it from growing too fast. In practice, `exp()` is usually enough.
@ -243,16 +245,19 @@ class StateDependentNoiseDistribution(Distribution):
this allows to ensure boundaries. this allows to ensure boundaries.
:param epsilon: (float) small value to avoid NaN due to numerical imprecision. :param epsilon: (float) small value to avoid NaN due to numerical imprecision.
""" """
def __init__(self, action_dim, use_expln=False, def __init__(self, action_dim, full_std=True, use_expln=False,
squash_output=False, epsilon=1e-6): squash_output=False, epsilon=1e-6):
super(StateDependentNoiseDistribution, self).__init__() super(StateDependentNoiseDistribution, self).__init__()
self.distribution = None self.distribution = None
self.action_dim = action_dim self.action_dim = action_dim
self.latent_dim = None
self.mean_actions = None self.mean_actions = None
self.log_std = None self.log_std = None
self.weights_dist = None self.weights_dist = None
self.exploration_mat = None self.exploration_mat = None
self.use_expln = use_expln self.use_expln = use_expln
self.full_std = full_std
self.epsilon = epsilon
if squash_output: if squash_output:
print("== Using TanhBijector ===") print("== Using TanhBijector ===")
self.bijector = TanhBijector(epsilon) self.bijector = TanhBijector(epsilon)
@ -271,12 +276,17 @@ class StateDependentNoiseDistribution(Distribution):
# From SDE paper, it allows to keep variance # From SDE paper, it allows to keep variance
# above zero and prevent it from growing too fast # above zero and prevent it from growing too fast
if log_std <= 0: if log_std <= 0:
return th.exp(log_std) std = th.exp(log_std)
else: else:
return th.log(log_std + 1.0) + 1.0 std = th.log(log_std + 1.0) + 1.0
else: else:
# Use normal exponential # Use normal exponential
return th.exp(log_std) std = th.exp(log_std)
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
def sample_weights(self, log_std): def sample_weights(self, log_std):
""" """
@ -285,11 +295,11 @@ class StateDependentNoiseDistribution(Distribution):
:param log_std: (th.Tensor) :param log_std: (th.Tensor)
""" """
# TODO: reduce the number of learned dimensions (cf TD3) std = self.get_std(log_std)
self.weights_dist = Normal(th.zeros_like(log_std), self.get_std(log_std)) self.weights_dist = Normal(th.zeros_like(std), std)
self.exploration_mat = self.weights_dist.rsample() self.exploration_mat = self.weights_dist.rsample()
def proba_distribution_net(self, latent_dim, log_std_init=0.0): def proba_distribution_net(self, latent_dim, log_std_init=-2.0):
""" """
Create the layers and parameter that represent the distribution: Create the layers and parameter that represent the distribution:
one output will be the deterministic action, the other parameter will be the one output will be the deterministic action, the other parameter will be the
@ -299,10 +309,17 @@ class StateDependentNoiseDistribution(Distribution):
:param log_std_init: (float) Initial value for the log standard deviation :param log_std_init: (float) Initial value for the log standard deviation
:return: (nn.Linear, nn.Parameter) :return: (nn.Linear, nn.Parameter)
""" """
mean_actions = nn.Linear(latent_dim, self.action_dim) # Network for the deterministic action, it represents the mean of the distribution
log_std = nn.Parameter(th.ones(latent_dim, self.action_dim) * log_std_init) mean_actions_net = nn.Linear(latent_dim, self.action_dim)
self.latent_dim = latent_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)
# 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) self.sample_weights(log_std)
return mean_actions, 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_pi, deterministic=False):
""" """
@ -314,7 +331,7 @@ class StateDependentNoiseDistribution(Distribution):
:return: (th.Tensor) :return: (th.Tensor)
""" """
variance = th.mm(latent_pi.detach() ** 2, self.get_std(log_std) ** 2) variance = th.mm(latent_pi.detach() ** 2, self.get_std(log_std) ** 2)
self.distribution = Normal(mean_actions, th.sqrt(variance)) self.distribution = Normal(mean_actions, th.sqrt(variance + self.epsilon))
if deterministic: if deterministic:
action = self.mode() action = self.mode()
@ -328,8 +345,11 @@ class StateDependentNoiseDistribution(Distribution):
return self.bijector.forward(action) return self.bijector.forward(action)
return action return action
def get_noise(self, latent_pi):
return th.mm(latent_pi.detach(), self.exploration_mat)
def sample(self, latent_pi): def sample(self, latent_pi):
noise = th.mm(latent_pi.detach(), self.exploration_mat) noise = self.get_noise(latent_pi)
action = self.distribution.mean + noise action = self.distribution.mean + noise
if self.bijector is not None: if self.bijector is not None:
return self.bijector.forward(action) return self.bijector.forward(action)
@ -405,26 +425,30 @@ class TanhBijector(object):
return th.log(1 - th.tanh(x) ** 2 + self.epsilon) return th.log(1 - th.tanh(x) ** 2 + self.epsilon)
def make_proba_distribution(action_space, use_sde=False): def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None):
""" """
Return an instance of Distribution for the correct type of action space Return an instance of Distribution for the correct type of action space
:param action_space: (Gym Space) the input action space :param action_space: (Gym Space) the input action space
:param use_sde: (bool) Force the use of StateDependentNoiseDistribution :param use_sde: (bool) Force the use of StateDependentNoiseDistribution
instead of DiagGaussianDistribution instead of DiagGaussianDistribution
:param dist_kwargs: (dict) Keyword arguments to pass to the probabilty distribution
:return: (Distribution) the approriate Distribution object :return: (Distribution) the approriate Distribution object
""" """
if dist_kwargs is None:
dist_kwargs = {}
if isinstance(action_space, spaces.Box): if isinstance(action_space, spaces.Box):
assert len(action_space.shape) == 1, "Error: the action space must be a vector" assert len(action_space.shape) == 1, "Error: the action space must be a vector"
if use_sde: if use_sde:
return StateDependentNoiseDistribution(action_space.shape[0]) return StateDependentNoiseDistribution(action_space.shape[0], **dist_kwargs)
return DiagGaussianDistribution(action_space.shape[0]) return DiagGaussianDistribution(action_space.shape[0], **dist_kwargs)
elif isinstance(action_space, spaces.Discrete): elif isinstance(action_space, spaces.Discrete):
return CategoricalDistribution(action_space.n) return CategoricalDistribution(action_space.n, **dist_kwargs)
# elif isinstance(action_space, spaces.MultiDiscrete): # elif isinstance(action_space, spaces.MultiDiscrete):
# return MultiCategoricalDistribution(action_space.nvec) # return MultiCategoricalDistribution(action_space.nvec, **dist_kwargs)
# elif isinstance(action_space, spaces.MultiBinary): # elif isinstance(action_space, spaces.MultiBinary):
# return BernoulliDistribution(action_space.n) # return BernoulliDistribution(action_space.n, **dist_kwargs)
else: else:
raise NotImplementedError("Error: probability distribution, not implemented for action space of type {}." raise NotImplementedError("Error: probability distribution, not implemented for action space of type {}."
.format(type(action_space)) + .format(type(action_space)) +

91
torchy_baselines/common/policies.py
View file

@ -1,3 +1,5 @@
from itertools import zip_longest
import torch as th import torch as th
import torch.nn as nn import torch.nn as nn
@ -141,3 +143,92 @@ def register_policy(name, policy):
raise ValueError("Error: the name {} is alreay registered for a different policy, will not override." raise ValueError("Error: the name {} is alreay registered for a different policy, will not override."
.format(name)) .format(name))
_policy_registry[sub_class][name] = policy _policy_registry[sub_class][name] = policy
class MlpExtractor(nn.Module):
"""
Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and
a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many
of them are shared between the policy network and the value network. It is assumed to be a list with the following
structure:
1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer.
If the number of ints is zero, there will be no shared layers.
2. An optional dict, to specify the following non-shared layers for the value network and the policy network.
It is formatted like ``dict(vf=[<value layer sizes>], pi=[<policy layer sizes>])``.
If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed.
For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value
network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec
would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128
would be specified as [128, 128].
Adapted from Stable Baselines.
:param feature_dim: (int) Dimension of the feature vector (can be the output of a CNN)
:param net_arch: ([int or dict]) The specification of the policy and value networks.
See above for details on its formatting.
:param activation_fn: (nn.Module) The activation function to use for the networks.
:param device: (th.device)
"""
def __init__(self, feature_dim, net_arch, activation_fn, device='cpu'):
super(MlpExtractor, self).__init__()
shared_net, policy_net, value_net = [], [], []
policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network
value_only_layers = [] # Layer sizes of the network that only belongs to the value network
last_layer_dim_shared = feature_dim
# Iterate through the shared layers and build the shared parts of the network
for idx, layer in enumerate(net_arch):
if isinstance(layer, int): # Check that this is a shared layer
layer_size = layer
# TODO: give layer a meaningful name
shared_net.append(nn.Linear(last_layer_dim_shared, layer_size))
shared_net.append(activation_fn())
last_layer_dim_shared = layer_size
else:
assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts"
if 'pi' in layer:
assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers."
policy_only_layers = layer['pi']
if 'vf' in layer:
assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers."
value_only_layers = layer['vf']
break # From here on the network splits up in policy and value network
last_layer_dim_pi = last_layer_dim_shared
last_layer_dim_vf = last_layer_dim_shared
# Build the non-shared part of the network
for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)):
if pi_layer_size is not None:
assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers."
policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size))
policy_net.append(activation_fn())
last_layer_dim_pi = pi_layer_size
if vf_layer_size is not None:
assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers."
value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size))
value_net.append(activation_fn())
last_layer_dim_vf = vf_layer_size
# Save dim, used to create the distributions
self.latent_dim_pi = last_layer_dim_pi
self.latent_dim_vf = last_layer_dim_vf
# Create networks
# If the list of layers is empty, the network will just act as an Identity module
self.shared_net = nn.Sequential(*shared_net).to(device)
self.policy_net = nn.Sequential(*policy_net).to(device)
self.value_net = nn.Sequential(*value_net).to(device)
def forward(self, features):
"""
:return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network.
If all layers are shared, then ``latent_policy == latent_value``
"""
shared_latent = self.shared_net(features)
return self.policy_net(shared_latent), self.value_net(shared_latent)

31
torchy_baselines/common/vec_env/__init__.py
View file

@ -1,7 +1,38 @@
# flake8: noqa F401 # flake8: noqa F401
from copy import deepcopy
from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError,\ from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError,\
VecEnv, VecEnvWrapper, CloudpickleWrapper VecEnv, VecEnvWrapper, CloudpickleWrapper
from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv
from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack
from torchy_baselines.common.vec_env.vec_normalize import VecNormalize from torchy_baselines.common.vec_env.vec_normalize import VecNormalize
def unwrap_vec_normalize(env):
"""
:param env: (gym.Env)
:return: (VecNormalize)
"""
env_tmp = env
while isinstance(env_tmp, VecEnvWrapper):
if isinstance(env_tmp, VecNormalize):
return env_tmp
env_tmp = env_tmp.venv
return None
# Define here to avoid circular import
def sync_envs_normalization(env, eval_env):
"""
Sync eval env and train env when using VecNormalize
:param env: (gym.Env)
:param eval_env: (gym.Env)
"""
env_tmp, eval_env_tmp = env, eval_env
while isinstance(env_tmp, VecEnvWrapper):
if isinstance(env_tmp, VecNormalize):
eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms)
env_tmp = env_tmp.venv
eval_env_tmp.venv

39
torchy_baselines/common/vec_env/vec_normalize.py
View file

@ -36,6 +36,7 @@ class VecNormalize(VecEnvWrapper):
self.norm_obs = norm_obs self.norm_obs = norm_obs
self.norm_reward = norm_reward self.norm_reward = norm_reward
self.old_obs = np.array([]) self.old_obs = np.array([])
self.old_reward = np.array([])
def step_wait(self): def step_wait(self):
""" """
@ -46,12 +47,13 @@ class VecNormalize(VecEnvWrapper):
""" """
obs, rews, news, infos = self.venv.step_wait() obs, rews, news, infos = self.venv.step_wait()
self.ret = self.ret * self.gamma + rews self.ret = self.ret * self.gamma + rews
self.old_obs = obs self.old_obs = obs.copy()
self.old_reward = rews.copy()
obs = self._normalize_observation(obs) obs = self._normalize_observation(obs)
if self.norm_reward: if self.norm_reward:
if self.training: if self.training:
self.ret_rms.update(self.ret) self.ret_rms.update(self.ret)
rews = np.clip(rews / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward) rews = self.normalize_reward(rews)
self.ret[news] = 0 self.ret[news] = 0
return obs, rews, news, infos return obs, rews, news, infos
@ -62,12 +64,31 @@ class VecNormalize(VecEnvWrapper):
if self.norm_obs: if self.norm_obs:
if self.training: if self.training:
self.obs_rms.update(obs) self.obs_rms.update(obs)
obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs, return self.normalize_obs(obs)
self.clip_obs)
return obs
else: else:
return obs return obs
def normalize_obs(self, obs):
if self.norm_obs:
return 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):
if self.norm_reward:
return np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
return reward
def unnormalize_obs(self, obs):
if self.norm_obs:
return (obs * np.sqrt(self.obs_rms.var + self.epsilon)) + self.obs_rms.mean
return obs
def unnormalize_reward(self, reward):
if self.norm_reward:
return reward * np.sqrt(self.ret_rms.var + self.epsilon)
return reward
def get_original_obs(self): def get_original_obs(self):
""" """
returns the unnormalized observation returns the unnormalized observation
@ -76,6 +97,14 @@ class VecNormalize(VecEnvWrapper):
""" """
return self.old_obs return self.old_obs
def get_original_reward(self):
"""
returns the unnormalized observation
:return: (numpy float)
"""
return self.old_reward
def reset(self): def reset(self):
""" """
Reset all environments Reset all environments

150
torchy_baselines/ppo/policies.py
View file

@ -1,113 +1,44 @@
from functools import partial from functools import partial
from itertools import zip_longest
import torch as th import torch as th
import torch.nn as nn import torch.nn as nn
import numpy as np import numpy as np
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor
from torchy_baselines.common.distributions import make_proba_distribution,\ from torchy_baselines.common.distributions import make_proba_distribution,\
DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
class MlpExtractor(nn.Module):
"""
Constructs an MLP that receives observations as an input and outputs a latent representation for the policy and
a value network. The ``net_arch`` parameter allows to specify the amount and size of the hidden layers and how many
of them are shared between the policy network and the value network. It is assumed to be a list with the following
structure:
1. An arbitrary length (zero allowed) number of integers each specifying the number of units in a shared layer.
If the number of ints is zero, there will be no shared layers.
2. An optional dict, to specify the following non-shared layers for the value network and the policy network.
It is formatted like ``dict(vf=[<value layer sizes>], pi=[<policy layer sizes>])``.
If it is missing any of the keys (pi or vf), no non-shared layers (empty list) is assumed.
For example to construct a network with one shared layer of size 55 followed by two non-shared layers for the value
network of size 255 and a single non-shared layer of size 128 for the policy network, the following layers_spec
would be used: ``[55, dict(vf=[255, 255], pi=[128])]``. A simple shared network topology with two layers of size 128
would be specified as [128, 128].
Adapted from Stable Baselines.
:param flat_observations: (th.Tensor) The observations to base policy and value function on.
:param net_arch: ([int or dict]) The specification of the policy and value networks.
See above for details on its formatting.
:param activation_fn: (nn.Module) The activation function to use for the networks.
:param device: (th.device)
"""
def __init__(self, feature_dim, net_arch, activation_fn, device='cpu'):
super(MlpExtractor, self).__init__()
shared_net, policy_net, value_net = [], [], []
policy_only_layers = [] # Layer sizes of the network that only belongs to the policy network
value_only_layers = [] # Layer sizes of the network that only belongs to the value network
last_layer_dim_shared = feature_dim
# Iterate through the shared layers and build the shared parts of the network
for idx, layer in enumerate(net_arch):
if isinstance(layer, int): # Check that this is a shared layer
layer_size = layer
# TODO: give layer a meaningful name
shared_net.append(nn.Linear(last_layer_dim_shared, layer_size))
shared_net.append(activation_fn())
last_layer_dim_shared = layer_size
else:
assert isinstance(layer, dict), "Error: the net_arch list can only contain ints and dicts"
if 'pi' in layer:
assert isinstance(layer['pi'], list), "Error: net_arch[-1]['pi'] must contain a list of integers."
policy_only_layers = layer['pi']
if 'vf' in layer:
assert isinstance(layer['vf'], list), "Error: net_arch[-1]['vf'] must contain a list of integers."
value_only_layers = layer['vf']
break # From here on the network splits up in policy and value network
last_layer_dim_pi = last_layer_dim_shared
last_layer_dim_vf = last_layer_dim_shared
# Build the non-shared part of the network
for idx, (pi_layer_size, vf_layer_size) in enumerate(zip_longest(policy_only_layers, value_only_layers)):
if pi_layer_size is not None:
assert isinstance(pi_layer_size, int), "Error: net_arch[-1]['pi'] must only contain integers."
policy_net.append(nn.Linear(last_layer_dim_pi, pi_layer_size))
policy_net.append(activation_fn())
last_layer_dim_pi = pi_layer_size
if vf_layer_size is not None:
assert isinstance(vf_layer_size, int), "Error: net_arch[-1]['vf'] must only contain integers."
value_net.append(nn.Linear(last_layer_dim_vf, vf_layer_size))
value_net.append(activation_fn())
last_layer_dim_vf = vf_layer_size
# Save dim, used to create the distributions
self.latent_dim_pi = last_layer_dim_pi
self.latent_dim_vf = last_layer_dim_vf
# Create networks
# If the list of layers is empty, the network will just act as an Identity module
self.shared_net = nn.Sequential(*shared_net).to(device)
self.policy_net = nn.Sequential(*policy_net).to(device)
self.value_net = nn.Sequential(*value_net).to(device)
def forward(self, features):
"""
:return: (th.Tensor, th.Tensor) latent_policy, latent_value of the specified network.
If all layers are shared, then ``latent_policy == latent_value``
"""
shared_latent = self.shared_net(features)
return self.policy_net(shared_latent), self.value_net(shared_latent)
class PPOPolicy(BasePolicy): 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 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 adam_epsilon: (float) Small values to avoid NaN in ADAM optimizer
:param ortho_init: (bool) Whether to use or not orthogonal initialization
: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
"""
def __init__(self, observation_space, action_space, def __init__(self, observation_space, action_space,
learning_rate, net_arch=None, device='cpu', learning_rate, net_arch=None, device='cpu',
activation_fn=nn.Tanh, adam_epsilon=1e-5, activation_fn=nn.Tanh, adam_epsilon=1e-5,
ortho_init=True, use_sde=False, log_std_init=0.0): ortho_init=True, use_sde=False,
log_std_init=0.0, full_std=True):
super(PPOPolicy, self).__init__(observation_space, action_space, device) super(PPOPolicy, self).__init__(observation_space, action_space, device)
self.obs_dim = self.observation_space.shape[0] self.obs_dim = self.observation_space.shape[0]
# Default network architecture, from stable-baselines
if net_arch is None: if net_arch is None:
net_arch = [dict(pi=[64], vf=[64])] net_arch = [dict(pi=[64, 64], vf=[64, 64])]
self.net_arch = net_arch self.net_arch = net_arch
self.activation_fn = activation_fn self.activation_fn = activation_fn
self.adam_epsilon = adam_epsilon self.adam_epsilon = adam_epsilon
@ -124,12 +55,24 @@ class PPOPolicy(BasePolicy):
self.features_extractor = nn.Flatten() self.features_extractor = nn.Flatten()
self.features_dim = self.obs_dim self.features_dim = self.obs_dim
self.log_std_init = log_std_init self.log_std_init = log_std_init
dist_kwargs = None
# Keyword arguments for SDE distribution
if use_sde:
dist_kwargs = {
'full_std': full_std,
'squash_output': False,
'use_expln': False
}
# Action distribution # Action distribution
self.action_dist = make_proba_distribution(action_space, use_sde=use_sde) self.action_dist = make_proba_distribution(action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
self._build(learning_rate) self._build(learning_rate)
def reset_noise_net(self): def reset_noise_net(self):
"""
Sample new weights for the exploration matrix.
"""
self.action_dist.sample_weights(self.log_std) self.action_dist.sample_weights(self.log_std)
def _build(self, learning_rate): def _build(self, learning_rate):
@ -147,7 +90,7 @@ class PPOPolicy(BasePolicy):
# with small initial weight for the output # with small initial weight for the output
if self.ortho_init: if self.ortho_init:
for module in [self.mlp_extractor, self.action_net, self.value_net]: for module in [self.mlp_extractor, self.action_net, self.value_net]:
# Values from stable-baselines check why # Values from stable-baselines, TODO: check why
gain = { gain = {
self.mlp_extractor: np.sqrt(2), self.mlp_extractor: np.sqrt(2),
self.action_net: 0.01, self.action_net: 0.01,
@ -185,15 +128,26 @@ class PPOPolicy(BasePolicy):
action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic) action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
return action.detach().cpu().numpy() return action.detach().cpu().numpy()
def get_policy_stats(self, obs, action, deterministic=False): def evaluate_actions(self, obs, action, deterministic=False):
"""
Evaluate actions according to the current policy,
given the observations.
:param obs: (th.Tensor)
:param action: (th.Tensor)
:param deterministic: (bool)
: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) latent_pi, latent_vf = self._get_latent(obs)
_, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic) _, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
log_prob = action_distribution.log_prob(action) log_prob = action_distribution.log_prob(action)
value = self.value_net(latent_vf) value = self.value_net(latent_vf)
return value, log_prob, action_distribution.entropy() return value, log_prob, action_distribution.entropy()
def value_forward(self): def value_forward(self, obs):
pass _, latent_vf = self._get_latent(obs)
return self.value_net(latent_vf)
MlpPolicy = PPOPolicy MlpPolicy = PPOPolicy

18
torchy_baselines/ppo/ppo.py
View file

@ -1,6 +1,5 @@
import os import os
import time import time
from copy import deepcopy
import gym import gym
from gym import spaces from gym import spaces
@ -18,7 +17,7 @@ from torchy_baselines.common.base_class import BaseRLModel
from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.common.evaluation import evaluate_policy
from torchy_baselines.common.buffers import RolloutBuffer from torchy_baselines.common.buffers import RolloutBuffer
from torchy_baselines.common.utils import explained_variance, get_schedule_fn from torchy_baselines.common.utils import explained_variance, get_schedule_fn
from torchy_baselines.common.vec_env import VecNormalize, VecEnvWrapper from torchy_baselines.common.vec_env import sync_envs_normalization
from torchy_baselines.common import logger from torchy_baselines.common import logger
from torchy_baselines.ppo.policies import PPOPolicy from torchy_baselines.ppo.policies import PPOPolicy
@ -205,7 +204,7 @@ class PPO(BaseRLModel):
# Convert discrete action for float to long # Convert discrete action for float to long
action = action.long().flatten() action = action.long().flatten()
values, log_prob, entropy = self.policy.get_policy_stats(obs, action) values, log_prob, entropy = self.policy.evaluate_actions(obs, action)
values = values.flatten() values = values.flatten()
# Normalize advantage # Normalize advantage
advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8) advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
@ -241,7 +240,8 @@ class PPO(BaseRLModel):
approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy()) approx_kl_divs.append(th.mean(old_log_prob - log_prob).detach().cpu().numpy())
if self.target_kl is not None and np.mean(approx_kl_divs) > 1.5 * self.target_kl: if self.target_kl is not None and np.mean(approx_kl_divs) > 1.5 * self.target_kl:
print("Early stopping at step {} due to reaching max kl: {:.2f}".format(it, np.mean(approx_kl_divs))) print("Early stopping at step {} due to reaching max kl: {:.2f}".format(gradient_step,
np.mean(approx_kl_divs)))
break break
explained_var = explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(), explained_var = explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),
@ -294,14 +294,8 @@ class PPO(BaseRLModel):
# Evaluate agent # Evaluate agent
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
timesteps_since_eval %= eval_freq timesteps_since_eval %= eval_freq
# TODO: move that to the base class sync_envs_normalization(self.env, eval_env)
# Sync eval env and train env when using VecNormalize
env_tmp, eval_env_tmp = self.env, eval_env
while isinstance(env_tmp, VecEnvWrapper):
if isinstance(env_tmp, VecNormalize):
eval_env_tmp.obs_rms = deepcopy(env_tmp.obs_rms)
env_tmp = env_tmp.venv
eval_env_tmp.venv
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes) mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
if self.tb_writer is not None: if self.tb_writer is not None:
self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps) self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)

9
torchy_baselines/sac/sac.py
View file

@ -8,6 +8,7 @@ from torchy_baselines.common.base_class import BaseRLModel
from torchy_baselines.common.buffers import ReplayBuffer from torchy_baselines.common.buffers import ReplayBuffer
from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.common.evaluation import evaluate_policy
from torchy_baselines.sac.policies import SACPolicy from torchy_baselines.sac.policies import SACPolicy
from torchy_baselines.common.vec_env import sync_envs_normalization
class SAC(BaseRLModel): class SAC(BaseRLModel):
@ -54,7 +55,7 @@ class SAC(BaseRLModel):
""" """
def __init__(self, policy, env, learning_rate=3e-4, buffer_size=int(1e6), def __init__(self, policy, env, learning_rate=3e-4, buffer_size=int(1e6),
learning_starts=100, batch_size=64, learning_starts=100, batch_size=256,
tau=0.005, ent_coef='auto', target_update_interval=1, tau=0.005, ent_coef='auto', target_update_interval=1,
train_freq=1, gradient_steps=1, n_episodes_rollout=-1, train_freq=1, gradient_steps=1, n_episodes_rollout=-1,
target_entropy='auto', action_noise=None, target_entropy='auto', action_noise=None,
@ -163,7 +164,7 @@ class SAC(BaseRLModel):
for gradient_step in range(gradient_steps): for gradient_step in range(gradient_steps):
# Sample replay buffer # Sample replay buffer
replay_data = self.replay_buffer.sample(batch_size) replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
obs, action_batch, next_obs, done, reward = replay_data obs, action_batch, next_obs, done, reward = replay_data
@ -257,9 +258,6 @@ class SAC(BaseRLModel):
self._update_current_progress(self.num_timesteps, total_timesteps) self._update_current_progress(self.num_timesteps, total_timesteps)
if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
if self.verbose > 1:
print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
self.num_timesteps, episode_num, episode_timesteps, episode_reward))
gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps
self.train(gradient_steps, batch_size=self.batch_size) self.train(gradient_steps, batch_size=self.batch_size)
@ -267,6 +265,7 @@ class SAC(BaseRLModel):
# Evaluate episode # Evaluate episode
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
timesteps_since_eval %= eval_freq timesteps_since_eval %= eval_freq
sync_envs_normalization(self.env, eval_env)
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes) mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
evaluations.append(mean_reward) evaluations.append(mean_reward)
if self.verbose > 0: if self.verbose > 0:

144
torchy_baselines/td3/policies.py
View file

@ -2,21 +2,120 @@ import torch as th
import torch.nn as nn 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
from torchy_baselines.common.distributions import StateDependentNoiseDistribution
class Actor(BaseNetwork): class Actor(BaseNetwork):
def __init__(self, obs_dim, action_dim, net_arch, activation_fn=nn.ReLU): """
Actor network (policy) for TD3.
: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 clip_noise: (float) Clip the magnitude of the noise
: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.
"""
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):
super(Actor, self).__init__() super(Actor, self).__init__()
# TODO: orthogonal initialization? self.latent_pi, self.log_std = None, None
actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True) self.weights_dist, self.exploration_mat = None, None
self.actor_net = nn.Sequential(*actor_net) self.use_sde, self.sde_optimizer = use_sde, None
self.action_dim = action_dim
self.full_std = full_std
def forward(self, obs): if use_sde:
return self.actor_net(obs) latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False)
self.latent_pi = nn.Sequential(*latent_pi)
# Create state dependent noise matrix (SDE)
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
squash_output=False)
action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=net_arch[-1],
log_std_init=log_std_init)
# Squash output
self.actor_net = 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)
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 _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 evaluate_actions(self, obs, action):
"""
Evaluate actions according to the current policy,
given the observations. Only useful when using SDE.
: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)
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.
"""
self.action_dist.sample_weights(self.log_std)
def forward(self, obs, deterministic=True):
if self.use_sde:
latent_pi = self.latent_pi(obs)
if deterministic:
return self.actor_net(latent_pi)
noise = self.action_dist.get_noise(latent_pi)
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
# action, _ = self._get_action_dist_from_latent(latent_pi)
# return action
else:
return self.actor_net(obs)
class Critic(BaseNetwork): class Critic(BaseNetwork):
"""
Critic network for TD3,
in fact it represents the action-state value function (Q-value function)
: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, def __init__(self, obs_dim, action_dim,
net_arch, activation_fn=nn.ReLU): net_arch, activation_fn=nn.ReLU):
super(Critic, self).__init__() super(Critic, self).__init__()
@ -38,11 +137,25 @@ class Critic(BaseNetwork):
class TD3Policy(BasePolicy): 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 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
"""
def __init__(self, observation_space, action_space, def __init__(self, observation_space, action_space,
learning_rate, net_arch=None, device='cpu', learning_rate, net_arch=None, device='cpu',
activation_fn=nn.ReLU): activation_fn=nn.ReLU, use_sde=False, log_std_init=-2,
clip_noise=None, lr_sde=3e-4):
super(TD3Policy, self).__init__(observation_space, action_space, device) super(TD3Policy, self).__init__(observation_space, action_space, device)
# Default network architecture, from the original paper
if net_arch is None: if net_arch is None:
net_arch = [400, 300] net_arch = [400, 300]
@ -56,8 +169,16 @@ class TD3Policy(BasePolicy):
'net_arch': self.net_arch, 'net_arch': self.net_arch,
'activation_fn': self.activation_fn '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
self.actor, self.actor_target = None, None self.actor, self.actor_target = None, None
self.critic, self.critic_target = None, None self.critic, self.critic_target = None, None
self.use_sde = use_sde
self.log_std_init = log_std_init
self._build(learning_rate) self._build(learning_rate)
def _build(self, learning_rate): def _build(self, learning_rate):
@ -71,14 +192,17 @@ class TD3Policy(BasePolicy):
self.critic_target.load_state_dict(self.critic.state_dict()) self.critic_target.load_state_dict(self.critic.state_dict())
self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1)) self.critic.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
def reset_noise(self):
return self.actor.reset_noise()
def make_actor(self): def make_actor(self):
return Actor(**self.net_args).to(self.device) return Actor(**self.actor_kwargs).to(self.device)
def make_critic(self): def make_critic(self):
return Critic(**self.net_args).to(self.device) return Critic(**self.net_args).to(self.device)
def forward(self, obs): def forward(self, obs, deterministic=True):
return self.actor(obs) return self.actor(obs, deterministic=deterministic)
MlpPolicy = TD3Policy MlpPolicy = TD3Policy

83
torchy_baselines/td3/td3.py
View file

@ -8,6 +8,7 @@ from torchy_baselines.common.base_class import BaseRLModel
from torchy_baselines.common.buffers import ReplayBuffer from torchy_baselines.common.buffers import ReplayBuffer
from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.common.evaluation import evaluate_policy
from torchy_baselines.td3.policies import TD3Policy from torchy_baselines.td3.policies import TD3Policy
from torchy_baselines.common.vec_env import sync_envs_normalization
class TD3(BaseRLModel): class TD3(BaseRLModel):
@ -37,6 +38,11 @@ class TD3(BaseRLModel):
:param target_policy_noise: (float) Standard deviation of gaussian noise added to target policy :param target_policy_noise: (float) Standard deviation of gaussian noise added to target policy
(smoothing noise) (smoothing noise)
:param target_noise_clip: (float) Limit for absolute value of target policy smoothing noise. :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_max_grad_norm: (float)
:param sde_ent_coef: (float)
:param sde_log_std_scheduler: (callable)
:param create_eval_env: (bool) Whether to create a second environment that will be :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) 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 :param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
@ -51,6 +57,7 @@ class TD3(BaseRLModel):
policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100,
train_freq=-1, gradient_steps=-1, n_episodes_rollout=1, 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, 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,
tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0, tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0,
seed=0, device='auto', _init_setup_model=True): seed=0, device='auto', _init_setup_model=True):
@ -58,7 +65,6 @@ class TD3(BaseRLModel):
create_eval_env=create_eval_env, seed=seed) create_eval_env=create_eval_env, seed=seed)
self.buffer_size = buffer_size self.buffer_size = buffer_size
# TODO: accept callables
self.learning_rate = learning_rate self.learning_rate = learning_rate
self.learning_starts = learning_starts self.learning_starts = learning_starts
self.train_freq = train_freq self.train_freq = train_freq
@ -72,6 +78,12 @@ class TD3(BaseRLModel):
self.target_noise_clip = target_noise_clip self.target_noise_clip = target_noise_clip
self.target_policy_noise = target_policy_noise 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
if _init_setup_model: if _init_setup_model:
self._setup_model() self._setup_model()
@ -81,7 +93,7 @@ class TD3(BaseRLModel):
self.set_random_seed(self.seed) self.set_random_seed(self.seed)
self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, 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.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.policy = self.policy.to(self.device)
self._create_aliases() self._create_aliases()
@ -91,12 +103,12 @@ class TD3(BaseRLModel):
self.critic = self.policy.critic self.critic = self.policy.critic
self.critic_target = self.policy.critic_target self.critic_target = self.policy.critic_target
def select_action(self, observation): def select_action(self, observation, deterministic=True):
# Normally not needed # Normally not needed
observation = np.array(observation) observation = np.array(observation)
with th.no_grad(): with th.no_grad():
observation = th.FloatTensor(observation.reshape(1, -1)).to(self.device) observation = th.FloatTensor(observation.reshape(1, -1)).to(self.device)
return self.actor(observation).cpu().numpy() return self.actor(observation, deterministic=deterministic).cpu().numpy()
def predict(self, observation, state=None, mask=None, deterministic=True): def predict(self, observation, state=None, mask=None, deterministic=True):
""" """
@ -108,7 +120,7 @@ class TD3(BaseRLModel):
:param deterministic: (bool) Whether or not to return deterministic actions. :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) :return: (np.ndarray, np.ndarray) the model's action and the next state (used in recurrent policies)
""" """
return self.unscale_action(self.select_action(observation)) return self.unscale_action(self.select_action(observation, deterministic=deterministic))
def train_critic(self, gradient_steps=1, batch_size=100, replay_data=None, tau=0.0): def train_critic(self, gradient_steps=1, batch_size=100, replay_data=None, tau=0.0):
# Update optimizer learning rate # Update optimizer learning rate
@ -117,7 +129,7 @@ class TD3(BaseRLModel):
for gradient_step in range(gradient_steps): for gradient_step in range(gradient_steps):
# Sample replay buffer # Sample replay buffer
if replay_data is None: if replay_data is None:
obs, action, next_obs, done, reward = self.replay_buffer.sample(batch_size) obs, action, next_obs, done, reward = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
else: else:
obs, action, next_obs, done, reward = replay_data obs, action, next_obs, done, reward = replay_data
@ -158,7 +170,7 @@ class TD3(BaseRLModel):
for gradient_step in range(gradient_steps): for gradient_step in range(gradient_steps):
# Sample replay buffer # Sample replay buffer
if replay_data is None: if replay_data is None:
obs, _, next_obs, done, reward = self.replay_buffer.sample(batch_size) obs, _, next_obs, done, reward = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
else: else:
obs, _, next_obs, done, reward = replay_data obs, _, next_obs, done, reward = replay_data
@ -183,13 +195,53 @@ class TD3(BaseRLModel):
for gradient_step in range(gradient_steps): for gradient_step in range(gradient_steps):
# Sample replay buffer # Sample replay buffer
replay_data = self.replay_buffer.sample(batch_size) replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
self.train_critic(replay_data=replay_data) self.train_critic(replay_data=replay_data)
# Delayed policy updates # Delayed policy updates
if gradient_step % policy_delay == 0: if gradient_step % policy_delay == 0:
self.train_actor(replay_data=replay_data, tau_actor=self.tau, tau_critic=self.tau) self.train_actor(replay_data=replay_data, tau_actor=self.tau, tau_critic=self.tau)
def train_sde(self):
# Update optimizer learning rate
# self._update_learning_rate(self.policy.optimizer)
# Unpack
obs, action, returns = [self.rollout_data[key] for key in ['observations', 'actions', 'returns']]
# TODO: avoid second computation of everything because of the gradient
log_prob, entropy = self.actor.evaluate_actions(obs, action)
# 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))
policy_loss = -(returns * log_prob).mean()
# Entropy loss favor exploration
entropy_loss = -th.mean(entropy)
loss = policy_loss + self.sde_ent_coef * entropy_loss
# Optimization step
self.actor.sde_optimizer.zero_grad()
loss.backward()
assert not th.isnan(log_prob).any(), log_prob
assert not th.isnan(entropy).any()
assert not th.isnan(self.actor.log_std.grad).any()
assert not th.isnan(self.actor.log_std).any()
# Clip grad norm
th.nn.utils.clip_grad_norm_([self.actor.log_std], self.sde_max_grad_norm)
self.actor.sde_optimizer.step()
del self.rollout_data
def learn(self, total_timesteps, callback=None, log_interval=4, def learn(self, total_timesteps, callback=None, log_interval=4,
eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True): eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True):
@ -219,9 +271,15 @@ class TD3(BaseRLModel):
self._update_current_progress(self.num_timesteps, total_timesteps) self._update_current_progress(self.num_timesteps, total_timesteps)
if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts: if self.num_timesteps > 0 and self.num_timesteps > self.learning_starts:
if self.verbose > 1:
print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format( if self.use_sde:
self.num_timesteps, episode_num, episode_timesteps, episode_reward)) if self.sde_log_std_scheduler is not None:
# Call the scheduler
value = self.sde_log_std_scheduler(self._current_progress)
self.actor.log_std.data = th.ones_like(self.actor.log_std) * value
else:
# On-policy gradient
self.train_sde()
gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps 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) self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay)
@ -229,6 +287,7 @@ class TD3(BaseRLModel):
# Evaluate episode # Evaluate episode
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
timesteps_since_eval %= eval_freq timesteps_since_eval %= eval_freq
sync_envs_normalization(self.env, eval_env)
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes) mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
evaluations.append(mean_reward) evaluations.append(mean_reward)
if self.verbose > 0: if self.verbose > 0:
@ -241,7 +300,7 @@ class TD3(BaseRLModel):
""" """
Returns a dict of all the optimizers and their parameters 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 {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict()} return {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict()}