mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-09-17 22:30:59 +00:00
Merge pull request #5 from Antonin-Raffin/feat/td3-sde
Off-Policy State Dependent Exploration
This commit is contained in:
Raffin, Antonin
GitHub Enterprise
commit
8874b9dd6b
19 changed files with 820 additions and 242 deletions
|
|
@ -21,6 +21,10 @@ TODO:
|
|||
- save/load
|
||||
- better predict
|
||||
- complete logger
|
||||
- SDE: learn the feature extractor?
|
||||
- Refactor: buffer with numpy array instead of pytorch
|
||||
- Refactor: remove duplicated code for evaluation
|
||||
- plotting? -> zoo
|
||||
|
||||
Later:
|
||||
- get_parameters / set_parameters
|
||||
|
|
|
|||
2
setup.py
2
setup.py
|
|
@ -34,7 +34,7 @@ setup(name='torchy_baselines',
|
|||
license="MIT",
|
||||
long_description="",
|
||||
long_description_content_type='text/markdown',
|
||||
version="0.0.5a",
|
||||
version="0.0.6a",
|
||||
)
|
||||
|
||||
# python setup.py sdist
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import gym
|
|||
import torch as th
|
||||
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.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,
|
||||
policy_kwargs=dict(log_std_init=0.0, ortho_init=False), seed=None)
|
||||
model.learn(total_timesteps=int(1000), log_interval=5, eval_freq=500, eval_env=eval_env)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_class", [TD3])
|
||||
def test_state_dependent_offpolicy_noise(model_class):
|
||||
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()
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
import gym
|
||||
import pytest
|
||||
import numpy as np
|
||||
|
||||
from torchy_baselines.common.running_mean_std import RunningMeanStd
|
||||
from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv
|
||||
from torchy_baselines.common.vec_env.vec_normalize import VecNormalize
|
||||
from torchy_baselines import CEMRL, SAC, TD3
|
||||
|
||||
ENV_ID = 'Pendulum-v0'
|
||||
|
||||
|
|
@ -39,3 +41,15 @@ def test_vec_env():
|
|||
actions = [env.action_space.sample()]
|
||||
obs, _, done, _ = env.step(actions)
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -4,4 +4,4 @@ from torchy_baselines.ppo import PPO
|
|||
from torchy_baselines.sac import SAC
|
||||
from torchy_baselines.td3 import TD3
|
||||
|
||||
__version__ = "0.0.5a"
|
||||
__version__ = "0.0.6a"
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
:param gamma: (float) Discount factor
|
||||
: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 vf_coef: (float) Value function coefficient for the loss calculation
|
||||
:param max_grad_norm: (float) The maximum value for the gradient clipping
|
||||
|
|
@ -92,7 +93,7 @@ class A2C(PPO):
|
|||
action = action.long().flatten()
|
||||
|
||||
# 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()
|
||||
|
||||
# Normalize advantage (not present in the original implementation)
|
||||
|
|
|
|||
|
|
@ -5,33 +5,42 @@ import numpy as np
|
|||
# or https://github.com/facebookresearch/nevergrad
|
||||
|
||||
class CEM(object):
|
||||
|
||||
"""
|
||||
Cross-entropy method with diagonal covariance (separable CEM)
|
||||
"""
|
||||
Cross-entropy method with diagonal covariance (separable CEM).
|
||||
|
||||
def __init__(self, num_params,
|
||||
mu_init=None,
|
||||
sigma_init=1e-3,
|
||||
pop_size=256,
|
||||
damp=1e-3,
|
||||
damp_limit=1e-5,
|
||||
parents=None,
|
||||
elitism=False,
|
||||
antithetic=False):
|
||||
:param num_params: (int) Number of parameters per individual (dimension of the problem)
|
||||
:param mu_init: (np.ndarray) Initial mean of the population distribution
|
||||
Taken to be zero if None is passed.
|
||||
: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 parents: (int) Number of parents used to compute the new distribution
|
||||
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__()
|
||||
# misc
|
||||
|
||||
self.num_params = num_params
|
||||
|
||||
# distribution parameters
|
||||
# Distribution parameters
|
||||
if mu_init is None:
|
||||
self.mu = np.zeros(self.num_params)
|
||||
else:
|
||||
self.mu = np.array(mu_init)
|
||||
|
||||
self.sigma = sigma_init
|
||||
self.damp = damp
|
||||
self.damp_limit = damp_limit
|
||||
# Damping parameters
|
||||
self.damping = damping_init
|
||||
self.damping_final = damping_final
|
||||
# Exponential moving average decay for damping
|
||||
self.tau = 0.95
|
||||
# Covariance matrix, here only the diagonal
|
||||
self.cov = self.sigma * np.ones(self.num_params)
|
||||
|
||||
# elite stuff
|
||||
|
|
@ -39,16 +48,20 @@ class CEM(object):
|
|||
self.elite = np.sqrt(self.sigma) * np.random.rand(self.num_params)
|
||||
self.elite_score = None
|
||||
|
||||
# sampling stuff
|
||||
# sampling parameters
|
||||
self.pop_size = pop_size
|
||||
self.antithetic = antithetic
|
||||
|
||||
if self.antithetic:
|
||||
assert (self.pop_size % 2 == 0), "Population size must be even"
|
||||
|
||||
if parents is None or parents <= 0:
|
||||
self.parents = pop_size // 2
|
||||
else:
|
||||
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)
|
||||
for i in range(1, self.parents + 1)])
|
||||
self.weights /= self.weights.sum()
|
||||
|
|
@ -56,43 +69,56 @@ class CEM(object):
|
|||
def ask(self, pop_size):
|
||||
"""
|
||||
Returns a list of candidates parameters
|
||||
|
||||
:param pop_size: (int)
|
||||
:return: ([np.ndarray])
|
||||
"""
|
||||
if self.antithetic and not pop_size % 2:
|
||||
epsilon_half = np.random.randn(pop_size // 2, self.num_params)
|
||||
epsilon = np.concatenate([epsilon_half, - epsilon_half])
|
||||
|
||||
else:
|
||||
epsilon = np.random.randn(pop_size, self.num_params)
|
||||
|
||||
inds = self.mu + epsilon * np.sqrt(self.cov)
|
||||
if self.elitism:
|
||||
inds[-1] = self.elite
|
||||
individuals = self.mu + epsilon * np.sqrt(self.cov)
|
||||
|
||||
return inds
|
||||
# Keep the best known individual in the population
|
||||
if self.elitism:
|
||||
individuals[-1] = self.elite
|
||||
|
||||
return individuals
|
||||
|
||||
def tell(self, solutions, scores):
|
||||
"""
|
||||
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 *= -1
|
||||
# Sort the individuals by fitness
|
||||
idx_sorted = np.argsort(scores)
|
||||
|
||||
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.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)
|
||||
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_score = scores[idx_sorted[0]]
|
||||
# print(self.cov)
|
||||
|
||||
def get_distrib_params(self):
|
||||
"""
|
||||
Returns the parameters of the distrubtion:
|
||||
the mean and sigma
|
||||
Returns the parameters of the distribution:
|
||||
the mean and standard deviation.
|
||||
|
||||
:return: (np.ndarray, np.ndarray)
|
||||
"""
|
||||
return np.copy(self.mu), np.copy(self.cov)
|
||||
|
|
|
|||
|
|
@ -5,21 +5,53 @@ import torch as th
|
|||
from torchy_baselines.cem_rl.cem import CEM
|
||||
from torchy_baselines.common.evaluation import evaluate_policy
|
||||
from torchy_baselines.td3.td3 import TD3
|
||||
from torchy_baselines.common.vec_env import sync_envs_normalization
|
||||
|
||||
|
||||
class CEMRL(TD3):
|
||||
"""
|
||||
Implementation of CEM-RL
|
||||
Implementation of CEM-RL, in fact CEM combined with TD3.
|
||||
|
||||
Paper: https://arxiv.org/abs/1810.01222
|
||||
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,
|
||||
damp=1e-3, damp_limit=1e-5, elitism=False, n_grad=5,
|
||||
policy_delay=2, batch_size=100,
|
||||
buffer_size=int(1e6), learning_rate=1e-3,
|
||||
action_noise=None, learning_starts=100, tau=0.005,
|
||||
damping_init=1e-3, damping_final=1e-5, elitism=False, n_grad=5,
|
||||
buffer_size=int(1e6), learning_rate=1e-3, policy_delay=2,
|
||||
learning_starts=100, gamma=0.99, batch_size=100, tau=0.005,
|
||||
action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5,
|
||||
n_episodes_rollout=1, update_style='original',
|
||||
tensorboard_log=None, create_eval_env=False,
|
||||
policy_kwargs=None, verbose=0, seed=0, device='auto',
|
||||
|
|
@ -27,18 +59,21 @@ class CEMRL(TD3):
|
|||
|
||||
super(CEMRL, self).__init__(policy, env,
|
||||
buffer_size=buffer_size, learning_rate=learning_rate, seed=seed, device=device,
|
||||
action_noise=action_noise, learning_starts=learning_starts,
|
||||
n_episodes_rollout=n_episodes_rollout, tau=tau,
|
||||
action_noise=action_noise, target_policy_noise=target_policy_noise,
|
||||
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_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)
|
||||
|
||||
# Evolution strategy method that follows cma-es interface (ask-tell)
|
||||
# for now, only CEM is implemented
|
||||
self.es = None
|
||||
self.sigma_init = sigma_init
|
||||
self.pop_size = pop_size
|
||||
self.damp = damp
|
||||
self.damp_limit = damp_limit
|
||||
self.damping_init = damping_init
|
||||
self.damping_final = damping_final
|
||||
self.elitism = elitism
|
||||
self.n_grad = n_grad
|
||||
self.es_params = None
|
||||
|
|
@ -52,7 +87,7 @@ class CEMRL(TD3):
|
|||
super(CEMRL, self)._setup_model()
|
||||
params_vector = self.actor.parameters_to_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,
|
||||
elitism=self.elitism)
|
||||
|
||||
|
|
@ -103,7 +138,7 @@ class CEMRL(TD3):
|
|||
n_training_steps = 2 * (actor_steps // self.n_grad)
|
||||
for it in range(n_training_steps):
|
||||
# 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)
|
||||
|
||||
# Delayed policy updates
|
||||
|
|
@ -118,6 +153,7 @@ class CEMRL(TD3):
|
|||
timesteps_since_eval %= eval_freq
|
||||
|
||||
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)
|
||||
evaluations.append(mean_reward)
|
||||
|
|
@ -150,10 +186,6 @@ class CEMRL(TD3):
|
|||
actor_steps += episode_timesteps
|
||||
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.es.tell(self.es_params, self.fitnesses)
|
||||
timesteps_since_eval += actor_steps
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import numpy as np
|
|||
|
||||
from torchy_baselines.common.policies import get_policy_from_name
|
||||
from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
|
||||
from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv
|
||||
from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize
|
||||
from torchy_baselines.common.monitor import Monitor
|
||||
from torchy_baselines.common import logger
|
||||
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 env: (Gym environment) The environment to learn from
|
||||
(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 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
|
||||
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
|
||||
or not in a Monitor wrapper.
|
||||
:param seed: (int) Seed for the pseudo random generators
|
||||
"""
|
||||
__metaclass__ = ABCMeta
|
||||
|
||||
|
|
@ -50,6 +56,8 @@ class BaseRLModel(object):
|
|||
print("Using {} device".format(self.device))
|
||||
|
||||
self.env = env
|
||||
# get VecNormalize object if needed
|
||||
self._vec_normalize_env = unwrap_vec_normalize(env)
|
||||
self.verbose = verbose
|
||||
self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs
|
||||
self.observation_space = None
|
||||
|
|
@ -60,10 +68,14 @@ class BaseRLModel(object):
|
|||
self.replay_buffer = None
|
||||
self.seed = seed
|
||||
self.action_noise = None
|
||||
# Used for SDE only
|
||||
self.rollout_data = None
|
||||
self.use_sde = False
|
||||
# Track the training progress (from 1 to 0)
|
||||
# this is used to update the learning rate
|
||||
self._current_progress = 1
|
||||
|
||||
# Create and wrap the env if needed
|
||||
if env is not None:
|
||||
if isinstance(env, str):
|
||||
if create_eval_env:
|
||||
|
|
@ -93,6 +105,12 @@ class BaseRLModel(object):
|
|||
" environment.")
|
||||
|
||||
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:
|
||||
eval_env = self.eval_env
|
||||
|
||||
|
|
@ -106,6 +124,9 @@ class BaseRLModel(object):
|
|||
"""
|
||||
Rescale the action from [low, high] to [-1, 1]
|
||||
(no need for symmetric action space)
|
||||
|
||||
:param action: (np.ndarray)
|
||||
:return: (np.ndarray)
|
||||
"""
|
||||
low, high = self.action_space.low, self.action_space.high
|
||||
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]
|
||||
(no need for symmetric action space)
|
||||
|
||||
:param scaled_action: (np.ndarray)
|
||||
:return: (np.ndarray)
|
||||
"""
|
||||
low, high = self.action_space.low, self.action_space.high
|
||||
return low + (0.5 * (scaled_action + 1.0) * (high - low))
|
||||
|
|
@ -126,7 +150,7 @@ class BaseRLModel(object):
|
|||
"""
|
||||
Compute current progress (from 1 to 0)
|
||||
|
||||
:param num_timesteps: (int)
|
||||
:param num_timesteps: (int) current number of timesteps
|
||||
:param total_timesteps: (int)
|
||||
"""
|
||||
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)
|
||||
|
||||
:return: (Gym Environment) The current environment
|
||||
:return: (gym.Env) The current environment
|
||||
"""
|
||||
return self.env
|
||||
|
||||
|
|
@ -190,7 +214,7 @@ class BaseRLModel(object):
|
|||
- observation_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:
|
||||
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.
|
||||
|
||||
: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.
|
||||
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 tb_log_name: (str) the name of the run for tensorboard log
|
||||
:param reset_num_timesteps: (bool) whether or not to reset the current timestep number (used in logging)
|
||||
:param eval_env: (gym.Env)
|
||||
:param eval_freq: (int)
|
||||
:param n_eval_episodes: (int)
|
||||
:return: (BaseRLModel) the trained model
|
||||
"""
|
||||
pass
|
||||
|
|
@ -407,21 +433,33 @@ class BaseRLModel(object):
|
|||
self.eval_env.seed(seed)
|
||||
|
||||
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.ep_info_buffer = deque(maxlen=100)
|
||||
|
||||
if self.action_noise is not None:
|
||||
self.action_noise.reset()
|
||||
|
||||
timesteps_since_eval, episode_num = 0, 0
|
||||
evaluations = []
|
||||
|
||||
if eval_env is not None and self.seed is not None:
|
||||
eval_env.seed(self.seed)
|
||||
|
||||
eval_env = self._get_eval_env(eval_env)
|
||||
obs = self.env.reset()
|
||||
return timesteps_since_eval, episode_num, evaluations, obs, eval_env
|
||||
|
||||
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])
|
||||
"""
|
||||
for info in infos:
|
||||
|
|
@ -434,13 +472,39 @@ class BaseRLModel(object):
|
|||
learning_starts=0, num_timesteps=0,
|
||||
replay_buffer=None, obs=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 = []
|
||||
total_timesteps = []
|
||||
total_steps, total_episodes = 0, 0
|
||||
assert isinstance(env, VecEnv)
|
||||
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:
|
||||
done = False
|
||||
# Reset environment: not needed for VecEnv
|
||||
|
|
@ -450,20 +514,29 @@ class BaseRLModel(object):
|
|||
while not done:
|
||||
# Select action randomly or according to policy
|
||||
if num_timesteps < learning_starts:
|
||||
action = np.array([self.action_space.sample()])
|
||||
# Warmup phase
|
||||
unscaled_action = np.array([self.action_space.sample()])
|
||||
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]
|
||||
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)
|
||||
if action_noise is not None:
|
||||
# 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
|
||||
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])]
|
||||
episode_reward += reward
|
||||
|
|
@ -473,14 +546,34 @@ class BaseRLModel(object):
|
|||
|
||||
# Store data in replay buffer
|
||||
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
|
||||
# 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
|
||||
episode_timesteps += 1
|
||||
total_steps += 1
|
||||
if n_steps > 0 and total_steps >= n_steps:
|
||||
if 0 < n_steps <= total_steps:
|
||||
break
|
||||
|
||||
if done:
|
||||
|
|
@ -495,29 +588,46 @@ class BaseRLModel(object):
|
|||
episode_num + total_episodes) % log_interval == 0:
|
||||
fps = int(num_timesteps / (time.time() - self.start_time))
|
||||
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:
|
||||
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("n_updates", n_updates)
|
||||
# logger.logkv("current_lr", current_lr)
|
||||
logger.logkv("fps", fps)
|
||||
logger.logkv('time_elapsed', int(time.time() - self.start_time))
|
||||
logger.logkv("total timesteps", num_timesteps)
|
||||
if self.use_sde:
|
||||
logger.logkv("std", (self.actor.get_std()).mean().item())
|
||||
logger.dumpkvs()
|
||||
|
||||
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
|
||||
|
||||
@staticmethod
|
||||
def _save_to_file_zip(save_path, data=None, params=None, opt_params=None):
|
||||
"""Save model to a zip archive
|
||||
|
||||
|
||||
:param save_path: (str) Where to store the model
|
||||
:param data: (dict) Class parameters being stored
|
||||
: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
|
||||
optimizer with its name and the state_dict
|
||||
optimizer with its name and the state_dict
|
||||
"""
|
||||
|
||||
# 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
|
||||
when saving the model.
|
||||
|
||||
|
||||
:return: ([str]) List of parameters that should be excluded from save
|
||||
"""
|
||||
return ["env", "eval_env", "replay_buffer", "rollout_buffer"]
|
||||
|
|
|
|||
|
|
@ -1,8 +1,19 @@
|
|||
import numpy as np
|
||||
import torch as th
|
||||
|
||||
from torchy_baselines.common.vec_env import unwrap_vec_normalize
|
||||
|
||||
|
||||
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):
|
||||
super(BaseBuffer, self).__init__()
|
||||
self.buffer_size = buffer_size
|
||||
|
|
@ -29,35 +40,68 @@ class BaseBuffer(object):
|
|||
return tensor.transpose(0, 1).reshape(shape[0] * shape[1], *shape[2:])
|
||||
|
||||
def size(self):
|
||||
"""
|
||||
:return: (int) The current size of the buffer
|
||||
"""
|
||||
if self.full:
|
||||
return self.buffer_size
|
||||
return self.pos
|
||||
|
||||
def get_pos(self):
|
||||
return self.pos
|
||||
|
||||
def add(self, *args, **kwargs):
|
||||
"""
|
||||
Add elements to the buffer.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def reset(self):
|
||||
"""
|
||||
Reset the buffer.
|
||||
"""
|
||||
self.pos = 0
|
||||
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
|
||||
batch_inds = th.LongTensor(
|
||||
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()
|
||||
|
||||
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):
|
||||
"""
|
||||
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):
|
||||
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.pos = 0
|
||||
|
||||
def _get_samples(self, batch_inds):
|
||||
return (self.observations[batch_inds, 0, :].to(self.device),
|
||||
def _get_samples(self, batch_inds, env=None):
|
||||
return (self._normalize_obs(self.observations[batch_inds, 0, :], env).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.rewards[batch_inds].to(self.device))
|
||||
self._normalize_reward(self.rewards[batch_inds], env).to(self.device))
|
||||
|
||||
|
||||
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',
|
||||
gae_lambda=1, gamma=0.99, n_envs=1):
|
||||
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):
|
||||
"""
|
||||
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 dones: ([bool])
|
||||
:param use_gae: (bool) Whether to use Generalized Advantage Estimation
|
||||
|
|
@ -151,6 +210,16 @@ class RolloutBuffer(BaseBuffer):
|
|||
self.advantages = self.returns - self.values
|
||||
|
||||
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:
|
||||
# Reshape 0-d tensor to avoid error
|
||||
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])
|
||||
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),
|
||||
self.actions[batch_inds].to(self.device),
|
||||
self.values[batch_inds].flatten().to(self.device),
|
||||
|
|
|
|||
|
|
@ -236,6 +236,8 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
compute the log probabilty of an action with that noise.
|
||||
|
||||
: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
|
||||
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.
|
||||
|
|
@ -243,16 +245,19 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
this allows to ensure boundaries.
|
||||
: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):
|
||||
super(StateDependentNoiseDistribution, self).__init__()
|
||||
self.distribution = None
|
||||
self.action_dim = action_dim
|
||||
self.latent_dim = None
|
||||
self.mean_actions = None
|
||||
self.log_std = None
|
||||
self.weights_dist = None
|
||||
self.exploration_mat = None
|
||||
self.use_expln = use_expln
|
||||
self.full_std = full_std
|
||||
self.epsilon = epsilon
|
||||
if squash_output:
|
||||
print("== Using TanhBijector ===")
|
||||
self.bijector = TanhBijector(epsilon)
|
||||
|
|
@ -271,12 +276,17 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
# From SDE paper, it allows to keep variance
|
||||
# above zero and prevent it from growing too fast
|
||||
if log_std <= 0:
|
||||
return th.exp(log_std)
|
||||
std = th.exp(log_std)
|
||||
else:
|
||||
return th.log(log_std + 1.0) + 1.0
|
||||
std = th.log(log_std + 1.0) + 1.0
|
||||
else:
|
||||
# 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):
|
||||
"""
|
||||
|
|
@ -285,11 +295,11 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
|
||||
:param log_std: (th.Tensor)
|
||||
"""
|
||||
# TODO: reduce the number of learned dimensions (cf TD3)
|
||||
self.weights_dist = Normal(th.zeros_like(log_std), self.get_std(log_std))
|
||||
std = self.get_std(log_std)
|
||||
self.weights_dist = Normal(th.zeros_like(std), std)
|
||||
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:
|
||||
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
|
||||
:return: (nn.Linear, nn.Parameter)
|
||||
"""
|
||||
mean_actions = nn.Linear(latent_dim, self.action_dim)
|
||||
log_std = nn.Parameter(th.ones(latent_dim, self.action_dim) * log_std_init)
|
||||
# Network for the deterministic action, it represents the mean of the distribution
|
||||
mean_actions_net = nn.Linear(latent_dim, self.action_dim)
|
||||
|
||||
self.latent_dim = latent_dim
|
||||
# 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)
|
||||
return mean_actions, log_std
|
||||
return mean_actions_net, log_std
|
||||
|
||||
def proba_distribution(self, mean_actions, log_std, latent_pi, deterministic=False):
|
||||
"""
|
||||
|
|
@ -314,7 +331,7 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
:return: (th.Tensor)
|
||||
"""
|
||||
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:
|
||||
action = self.mode()
|
||||
|
|
@ -328,8 +345,11 @@ class StateDependentNoiseDistribution(Distribution):
|
|||
return self.bijector.forward(action)
|
||||
return action
|
||||
|
||||
def get_noise(self, latent_pi):
|
||||
return th.mm(latent_pi.detach(), self.exploration_mat)
|
||||
|
||||
def sample(self, latent_pi):
|
||||
noise = th.mm(latent_pi.detach(), self.exploration_mat)
|
||||
noise = self.get_noise(latent_pi)
|
||||
action = self.distribution.mean + noise
|
||||
if self.bijector is not None:
|
||||
return self.bijector.forward(action)
|
||||
|
|
@ -405,26 +425,30 @@ class TanhBijector(object):
|
|||
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
|
||||
|
||||
:param action_space: (Gym Space) the input action space
|
||||
:param use_sde: (bool) Force the use of StateDependentNoiseDistribution
|
||||
instead of DiagGaussianDistribution
|
||||
:param dist_kwargs: (dict) Keyword arguments to pass to the probabilty distribution
|
||||
:return: (Distribution) the approriate Distribution object
|
||||
"""
|
||||
if dist_kwargs is None:
|
||||
dist_kwargs = {}
|
||||
|
||||
if isinstance(action_space, spaces.Box):
|
||||
assert len(action_space.shape) == 1, "Error: the action space must be a vector"
|
||||
if use_sde:
|
||||
return StateDependentNoiseDistribution(action_space.shape[0])
|
||||
return DiagGaussianDistribution(action_space.shape[0])
|
||||
return StateDependentNoiseDistribution(action_space.shape[0], **dist_kwargs)
|
||||
return DiagGaussianDistribution(action_space.shape[0], **dist_kwargs)
|
||||
elif isinstance(action_space, spaces.Discrete):
|
||||
return CategoricalDistribution(action_space.n)
|
||||
return CategoricalDistribution(action_space.n, **dist_kwargs)
|
||||
# elif isinstance(action_space, spaces.MultiDiscrete):
|
||||
# return MultiCategoricalDistribution(action_space.nvec)
|
||||
# return MultiCategoricalDistribution(action_space.nvec, **dist_kwargs)
|
||||
# elif isinstance(action_space, spaces.MultiBinary):
|
||||
# return BernoulliDistribution(action_space.n)
|
||||
# return BernoulliDistribution(action_space.n, **dist_kwargs)
|
||||
else:
|
||||
raise NotImplementedError("Error: probability distribution, not implemented for action space of type {}."
|
||||
.format(type(action_space)) +
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
from itertools import zip_longest
|
||||
|
||||
import torch as th
|
||||
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."
|
||||
.format(name))
|
||||
_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)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,38 @@
|
|||
# flake8: noqa F401
|
||||
from copy import deepcopy
|
||||
|
||||
from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError,\
|
||||
VecEnv, VecEnvWrapper, CloudpickleWrapper
|
||||
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.vec_frame_stack import VecFrameStack
|
||||
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
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ class VecNormalize(VecEnvWrapper):
|
|||
self.norm_obs = norm_obs
|
||||
self.norm_reward = norm_reward
|
||||
self.old_obs = np.array([])
|
||||
self.old_reward = np.array([])
|
||||
|
||||
def step_wait(self):
|
||||
"""
|
||||
|
|
@ -46,12 +47,13 @@ class VecNormalize(VecEnvWrapper):
|
|||
"""
|
||||
obs, rews, news, infos = self.venv.step_wait()
|
||||
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)
|
||||
if self.norm_reward:
|
||||
if self.training:
|
||||
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
|
||||
return obs, rews, news, infos
|
||||
|
||||
|
|
@ -62,12 +64,31 @@ class VecNormalize(VecEnvWrapper):
|
|||
if self.norm_obs:
|
||||
if self.training:
|
||||
self.obs_rms.update(obs)
|
||||
obs = np.clip((obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.epsilon), -self.clip_obs,
|
||||
self.clip_obs)
|
||||
return obs
|
||||
return self.normalize_obs(obs)
|
||||
else:
|
||||
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):
|
||||
"""
|
||||
returns the unnormalized observation
|
||||
|
|
@ -76,6 +97,14 @@ class VecNormalize(VecEnvWrapper):
|
|||
"""
|
||||
return self.old_obs
|
||||
|
||||
def get_original_reward(self):
|
||||
"""
|
||||
returns the unnormalized observation
|
||||
|
||||
:return: (numpy float)
|
||||
"""
|
||||
return self.old_reward
|
||||
|
||||
def reset(self):
|
||||
"""
|
||||
Reset all environments
|
||||
|
|
|
|||
|
|
@ -1,113 +1,44 @@
|
|||
from functools import partial
|
||||
from itertools import zip_longest
|
||||
|
||||
import torch as th
|
||||
import torch.nn as nn
|
||||
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,\
|
||||
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):
|
||||
"""
|
||||
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,
|
||||
learning_rate, net_arch=None, device='cpu',
|
||||
activation_fn=nn.Tanh, adam_epsilon=1e-5,
|
||||
ortho_init=True, use_sde=False, log_std_init=0.0):
|
||||
ortho_init=True, use_sde=False,
|
||||
log_std_init=0.0, full_std=True):
|
||||
super(PPOPolicy, self).__init__(observation_space, action_space, device)
|
||||
self.obs_dim = self.observation_space.shape[0]
|
||||
|
||||
|
||||
# Default network architecture, from stable-baselines
|
||||
if net_arch is None:
|
||||
net_arch = [dict(pi=[64], vf=[64])]
|
||||
net_arch = [dict(pi=[64, 64], vf=[64, 64])]
|
||||
|
||||
self.net_arch = net_arch
|
||||
self.activation_fn = activation_fn
|
||||
self.adam_epsilon = adam_epsilon
|
||||
|
|
@ -124,12 +55,24 @@ class PPOPolicy(BasePolicy):
|
|||
self.features_extractor = nn.Flatten()
|
||||
self.features_dim = self.obs_dim
|
||||
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
|
||||
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)
|
||||
|
||||
def reset_noise_net(self):
|
||||
"""
|
||||
Sample new weights for the exploration matrix.
|
||||
"""
|
||||
self.action_dist.sample_weights(self.log_std)
|
||||
|
||||
def _build(self, learning_rate):
|
||||
|
|
@ -147,7 +90,7 @@ class PPOPolicy(BasePolicy):
|
|||
# with small initial weight for the output
|
||||
if self.ortho_init:
|
||||
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 = {
|
||||
self.mlp_extractor: np.sqrt(2),
|
||||
self.action_net: 0.01,
|
||||
|
|
@ -185,15 +128,26 @@ class PPOPolicy(BasePolicy):
|
|||
action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
|
||||
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)
|
||||
_, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
|
||||
log_prob = action_distribution.log_prob(action)
|
||||
value = self.value_net(latent_vf)
|
||||
return value, log_prob, action_distribution.entropy()
|
||||
|
||||
def value_forward(self):
|
||||
pass
|
||||
def value_forward(self, obs):
|
||||
_, latent_vf = self._get_latent(obs)
|
||||
return self.value_net(latent_vf)
|
||||
|
||||
|
||||
MlpPolicy = PPOPolicy
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
import os
|
||||
import time
|
||||
from copy import deepcopy
|
||||
|
||||
import gym
|
||||
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.buffers import RolloutBuffer
|
||||
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.ppo.policies import PPOPolicy
|
||||
|
||||
|
|
@ -205,7 +204,7 @@ class PPO(BaseRLModel):
|
|||
# Convert discrete action for float to long
|
||||
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()
|
||||
# Normalize advantage
|
||||
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())
|
||||
|
||||
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
|
||||
|
||||
explained_var = explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),
|
||||
|
|
@ -294,14 +294,8 @@ class PPO(BaseRLModel):
|
|||
# Evaluate agent
|
||||
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
|
||||
timesteps_since_eval %= eval_freq
|
||||
# TODO: move that to the base class
|
||||
# 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
|
||||
sync_envs_normalization(self.env, eval_env)
|
||||
|
||||
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
|
||||
if self.tb_writer is not None:
|
||||
self.tb_writer.add_scalar('Eval/reward', mean_reward, self.num_timesteps)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from torchy_baselines.common.base_class import BaseRLModel
|
|||
from torchy_baselines.common.buffers import ReplayBuffer
|
||||
from torchy_baselines.common.evaluation import evaluate_policy
|
||||
from torchy_baselines.sac.policies import SACPolicy
|
||||
from torchy_baselines.common.vec_env import sync_envs_normalization
|
||||
|
||||
|
||||
class SAC(BaseRLModel):
|
||||
|
|
@ -54,7 +55,7 @@ class SAC(BaseRLModel):
|
|||
"""
|
||||
|
||||
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,
|
||||
train_freq=1, gradient_steps=1, n_episodes_rollout=-1,
|
||||
target_entropy='auto', action_noise=None,
|
||||
|
|
@ -163,7 +164,7 @@ class SAC(BaseRLModel):
|
|||
|
||||
for gradient_step in range(gradient_steps):
|
||||
# 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
|
||||
|
||||
|
|
@ -257,9 +258,6 @@ class SAC(BaseRLModel):
|
|||
self._update_current_progress(self.num_timesteps, total_timesteps)
|
||||
|
||||
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
|
||||
|
||||
self.train(gradient_steps, batch_size=self.batch_size)
|
||||
|
|
@ -267,6 +265,7 @@ class SAC(BaseRLModel):
|
|||
# Evaluate episode
|
||||
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
|
||||
timesteps_since_eval %= eval_freq
|
||||
sync_envs_normalization(self.env, eval_env)
|
||||
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
|
||||
evaluations.append(mean_reward)
|
||||
if self.verbose > 0:
|
||||
|
|
|
|||
|
|
@ -2,21 +2,120 @@ import torch as th
|
|||
import torch.nn as nn
|
||||
|
||||
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork
|
||||
from torchy_baselines.common.distributions import StateDependentNoiseDistribution
|
||||
|
||||
|
||||
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__()
|
||||
|
||||
# TODO: orthogonal initialization?
|
||||
actor_net = create_mlp(obs_dim, action_dim, net_arch, activation_fn, squash_out=True)
|
||||
self.actor_net = nn.Sequential(*actor_net)
|
||||
self.latent_pi, self.log_std = None, None
|
||||
self.weights_dist, self.exploration_mat = None, None
|
||||
self.use_sde, self.sde_optimizer = use_sde, None
|
||||
self.action_dim = action_dim
|
||||
self.full_std = full_std
|
||||
|
||||
def forward(self, obs):
|
||||
return self.actor_net(obs)
|
||||
if use_sde:
|
||||
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):
|
||||
"""
|
||||
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,
|
||||
net_arch, activation_fn=nn.ReLU):
|
||||
super(Critic, self).__init__()
|
||||
|
|
@ -38,11 +137,25 @@ class Critic(BaseNetwork):
|
|||
|
||||
|
||||
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,
|
||||
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)
|
||||
|
||||
# Default network architecture, from the original paper
|
||||
if net_arch is None:
|
||||
net_arch = [400, 300]
|
||||
|
||||
|
|
@ -56,8 +169,16 @@ class TD3Policy(BasePolicy):
|
|||
'net_arch': self.net_arch,
|
||||
'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.critic, self.critic_target = None, None
|
||||
self.use_sde = use_sde
|
||||
self.log_std_init = log_std_init
|
||||
self._build(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.optimizer = th.optim.Adam(self.critic.parameters(), lr=learning_rate(1))
|
||||
|
||||
def reset_noise(self):
|
||||
return self.actor.reset_noise()
|
||||
|
||||
def make_actor(self):
|
||||
return Actor(**self.net_args).to(self.device)
|
||||
return Actor(**self.actor_kwargs).to(self.device)
|
||||
|
||||
def make_critic(self):
|
||||
return Critic(**self.net_args).to(self.device)
|
||||
|
||||
def forward(self, obs):
|
||||
return self.actor(obs)
|
||||
def forward(self, obs, deterministic=True):
|
||||
return self.actor(obs, deterministic=deterministic)
|
||||
|
||||
|
||||
MlpPolicy = TD3Policy
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from torchy_baselines.common.base_class import BaseRLModel
|
|||
from torchy_baselines.common.buffers import ReplayBuffer
|
||||
from torchy_baselines.common.evaluation import evaluate_policy
|
||||
from torchy_baselines.td3.policies import TD3Policy
|
||||
from torchy_baselines.common.vec_env import sync_envs_normalization
|
||||
|
||||
|
||||
class TD3(BaseRLModel):
|
||||
|
|
@ -37,6 +38,11 @@ class TD3(BaseRLModel):
|
|||
: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 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
|
||||
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
|
||||
|
|
@ -51,6 +57,7 @@ class TD3(BaseRLModel):
|
|||
policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100,
|
||||
train_freq=-1, gradient_steps=-1, n_episodes_rollout=1,
|
||||
tau=0.005, action_noise=None, target_policy_noise=0.2, target_noise_clip=0.5,
|
||||
use_sde=False, sde_max_grad_norm=1, sde_ent_coef=0.0, sde_log_std_scheduler=None,
|
||||
tensorboard_log=None, create_eval_env=False, policy_kwargs=None, verbose=0,
|
||||
seed=0, device='auto', _init_setup_model=True):
|
||||
|
||||
|
|
@ -58,7 +65,6 @@ class TD3(BaseRLModel):
|
|||
create_eval_env=create_eval_env, seed=seed)
|
||||
|
||||
self.buffer_size = buffer_size
|
||||
# TODO: accept callables
|
||||
self.learning_rate = learning_rate
|
||||
self.learning_starts = learning_starts
|
||||
self.train_freq = train_freq
|
||||
|
|
@ -72,6 +78,12 @@ class TD3(BaseRLModel):
|
|||
self.target_noise_clip = target_noise_clip
|
||||
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:
|
||||
self._setup_model()
|
||||
|
||||
|
|
@ -81,7 +93,7 @@ class TD3(BaseRLModel):
|
|||
self.set_random_seed(self.seed)
|
||||
self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
|
||||
self.policy = self.policy_class(self.observation_space, self.action_space,
|
||||
self.learning_rate, device=self.device, **self.policy_kwargs)
|
||||
self.learning_rate, use_sde=self.use_sde, device=self.device, **self.policy_kwargs)
|
||||
self.policy = self.policy.to(self.device)
|
||||
self._create_aliases()
|
||||
|
||||
|
|
@ -91,12 +103,12 @@ class TD3(BaseRLModel):
|
|||
self.critic = self.policy.critic
|
||||
self.critic_target = self.policy.critic_target
|
||||
|
||||
def select_action(self, observation):
|
||||
def select_action(self, observation, deterministic=True):
|
||||
# Normally not needed
|
||||
observation = np.array(observation)
|
||||
with th.no_grad():
|
||||
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):
|
||||
"""
|
||||
|
|
@ -108,7 +120,7 @@ class TD3(BaseRLModel):
|
|||
: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 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):
|
||||
# Update optimizer learning rate
|
||||
|
|
@ -117,7 +129,7 @@ class TD3(BaseRLModel):
|
|||
for gradient_step in range(gradient_steps):
|
||||
# Sample replay buffer
|
||||
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:
|
||||
obs, action, next_obs, done, reward = replay_data
|
||||
|
||||
|
|
@ -158,7 +170,7 @@ class TD3(BaseRLModel):
|
|||
for gradient_step in range(gradient_steps):
|
||||
# Sample replay buffer
|
||||
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:
|
||||
obs, _, next_obs, done, reward = replay_data
|
||||
|
||||
|
|
@ -183,13 +195,53 @@ class TD3(BaseRLModel):
|
|||
for gradient_step in range(gradient_steps):
|
||||
|
||||
# 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)
|
||||
|
||||
# Delayed policy updates
|
||||
if gradient_step % policy_delay == 0:
|
||||
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,
|
||||
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)
|
||||
|
||||
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))
|
||||
|
||||
if self.use_sde:
|
||||
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
|
||||
self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay)
|
||||
|
|
@ -229,6 +287,7 @@ class TD3(BaseRLModel):
|
|||
# Evaluate episode
|
||||
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
|
||||
timesteps_since_eval %= eval_freq
|
||||
sync_envs_normalization(self.env, eval_env)
|
||||
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
|
||||
evaluations.append(mean_reward)
|
||||
if self.verbose > 0:
|
||||
|
|
@ -241,7 +300,7 @@ class TD3(BaseRLModel):
|
|||
"""
|
||||
Returns a dict of all the optimizers and their parameters
|
||||
|
||||
:return: (Dict) of optimizer names and their state_dict
|
||||
:return: (Dict) of optimizer names and their state_dict
|
||||
"""
|
||||
return {"actor": self.actor.optimizer.state_dict(), "critic": self.critic.optimizer.state_dict()}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue