saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-23 19:32:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add eval env and clip vf

Browse source
This commit is contained in:
Antonin RAFFIN 2019-09-19 17:18:41 +02:00
parent fe8b415cbf
commit cc4380eccd
6 changed files with 55 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
torchy_baselines/cem_rl/cem_rl.py
View file

@ -52,7 +52,7 @@ class CEMRL(TD3):
elitism=self.elitism)
def learn(self, total_timesteps, callback=None, log_interval=100,
eval_freq=-1, n_eval_episodes=5, tb_log_name="CEMRL", reset_num_timesteps=True):
eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="CEMRL", reset_num_timesteps=True):
timesteps_since_eval, actor_steps = 0, 0
episode_num = 0
@ -110,13 +110,13 @@ class CEMRL(TD3):
# Get the params back in the population
self.es_params[i] = self.actor.parameters_to_vector()
# Evaluate episode
if 0 < eval_freq <= timesteps_since_eval:
# Evaluate agent
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
timesteps_since_eval %= eval_freq
self.actor.load_from_vector(self.es.mu)
mean_reward, _ = evaluate_policy(self, self.env, n_eval_episodes)
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
evaluations.append(mean_reward)
if self.verbose > 0:

2
torchy_baselines/common/base_class.py
View file

@ -110,7 +110,7 @@ class BaseRLModel(object):
@abstractmethod
def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="run",
reset_num_timesteps=True):
eval_env=None, eval_freq=-1, n_eval_episodes=5, reset_num_timesteps=True):
"""
Return a trained model.

1
torchy_baselines/common/buffers.py
View file

@ -125,6 +125,7 @@ class RolloutBuffer(BaseBuffer):
def _get_samples(self, batch_inds):
return (self.states[batch_inds].to(self.device),
self.actions[batch_inds].to(self.device),
self.values[batch_inds].flatten().to(self.device),
self.log_probs[batch_inds].flatten().to(self.device),
self.advantages[batch_inds].flatten().to(self.device),
self.returns[batch_inds].flatten().to(self.device))

5
torchy_baselines/ppo/policies.py
View file

@ -8,7 +8,7 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create
class PPOPolicy(BasePolicy):
def __init__(self, observation_space, action_space,
learning_rate=1e-3, net_arch=None, device='cpu',
activation_fn=nn.Tanh):
activation_fn=nn.Tanh, adam_epsilon=1e-5):
super(PPOPolicy, self).__init__(observation_space, action_space, device)
self.state_dim = self.observation_space.shape[0]
self.action_dim = self.action_space.shape[0]
@ -16,6 +16,7 @@ class PPOPolicy(BasePolicy):
net_arch = [64, 64]
self.net_arch = net_arch
self.activation_fn = activation_fn
self.adam_epsilon = adam_epsilon
self.net_args = {
'input_dim': self.state_dim,
'output_dim': -1,
@ -41,7 +42,7 @@ class PPOPolicy(BasePolicy):
for module in [self.shared_net, self.actor_net, self.value_net]:
module.apply(self.init_weights)
self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate)
self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate, eps=self.adam_epsilon)
def forward(self, state, deterministic=False):
state = th.FloatTensor(state).to(self.device)

62
torchy_baselines/ppo/ppo.py
View file

@ -24,7 +24,8 @@ class PPO(BaseRLModel):
learning_rate=3e-4, seed=0, device='auto',
n_optim=5, batch_size=64, n_steps=256,
gamma=0.99, lambda_=0.95, clip_range=0.2,
ent_coef=0.01, vf_coef=0.5,
ent_coef=0.01, vf_coef=0.5, max_grad_norm=0.5,
target_kl=None, clip_range_vf=None,
_init_setup_model=True):
super(PPO, self).__init__(policy, env, PPOPolicy, policy_kwargs, verbose, device)
@ -40,7 +41,10 @@ class PPO(BaseRLModel):
self.clip_range = clip_range
self.ent_coef = ent_coef
self.vf_coef = vf_coef
self.max_grad_norm = max_grad_norm
self.rollout_buffer = None
self.target_kl = target_kl
self.clip_range_vf = clip_range_vf
if _init_setup_model:
self._setup_model()
@ -91,6 +95,7 @@ class PPO(BaseRLModel):
action = action.flatten().cpu().numpy()
# Rescale and perform action
# TODO: clip only when using Box action space
new_obs, reward, done, _ = env.step(np.clip(action, -self.max_action, self.max_action))
n_steps += 1
@ -108,41 +113,61 @@ class PPO(BaseRLModel):
# TODO: replace with iterator?
for it in range(n_iterations):
approx_kl_divs = []
# Sample replay buffer
for replay_data in self.rollout_buffer.get(batch_size):
# Unpack
state, action, old_log_prob, advantage, return_batch = replay_data
state, action, old_values, old_log_prob, advantage, return_batch = replay_data
values, log_prob, entropy = self.policy.get_policy_stats(state, action)
values = values.flatten()
# Normalize advantage
# advs = returns - values
advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
# ratio between old and new policy, should be one at the first iteration
ratio = th.exp(log_prob - old_log_prob)
# clipped surrogate loss
policy_loss_1 = advantage * ratio
policy_loss_2 = advantage * th.clamp(ratio, 1 - self.clip_range, 1 + self.clip_range)
policy_loss = -th.min(policy_loss_1, policy_loss_2).mean()
# value_loss = th.mean((return_batch - value)**2)
value_loss = F.mse_loss(return_batch, values.flatten())
if self.clip_range_vf is None:
# No clipping
values_pred = values
else:
# Clip the different between old and new value
# NOTE: this depends on the reward scaling
values_pred = old_values + th.clamp(values - old_values, -self.clip_range_vf, self.clip_range_vf)
# Value loss using the TD(lambda_) target
value_loss = F.mse_loss(return_batch, values_pred)
# Entropy loss favor exploration
entropy_loss = th.mean(entropy)
loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss
# loss = policy_loss
# TODO: check kl div
# self.approxkl = .5 * tf.reduce_mean(tf.square(neglogpac - self.old_neglog_pac_ph))
# approx_kl_div = th.mean(old_log_prob - log_prob)
# Optimization step
self.policy.optimizer.zero_grad()
loss.backward()
# TODO: clip grad norm?
# nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
# Clip grad norm
th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm)
self.policy.optimizer.step()
# print(value_loss.item())
# print(explained_variance(return_batch.numpy(), values.flatten().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:
print("Early stopping at step {} due to reaching max kl: {:.2f}".format(it, np.mean(approx_kl_divs)))
break
# print(explained_variance(self.rollout_buffer.returns.flatten().cpu().numpy(),
# self.rollout_buffer.values.flatten().cpu().numpy()))
def learn(self, total_timesteps, callback=None, log_interval=100,
eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", reset_num_timesteps=True):
eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", reset_num_timesteps=True):
timesteps_since_eval = 0
episode_num = 0
@ -157,19 +182,18 @@ class PPO(BaseRLModel):
if callback(locals(), globals()) is False:
break
# TODO: avoid reset using obs=obs and test env
obs = self.collect_rollouts(self.env, self.rollout_buffer, n_rollout_steps=self.n_steps,
obs=None)
obs=obs)
episode_num += 1
self.num_timesteps += self.n_steps
timesteps_since_eval += self.n_steps
self.train(self.n_optim, batch_size=self.batch_size)
# Evaluate episode
if 0 < eval_freq <= timesteps_since_eval:
# Evaluate agent
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
timesteps_since_eval %= eval_freq
mean_reward, _ = evaluate_policy(self, self.env, n_eval_episodes)
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
evaluations.append(mean_reward)
if self.verbose > 0:
print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))

6
torchy_baselines/td3/td3.py
View file

@ -147,7 +147,7 @@ class TD3(BaseRLModel):
self.train_actor(replay_data=replay_data)
def learn(self, total_timesteps, callback=None, log_interval=100,
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):
timesteps_since_eval = 0
episode_num = 0
@ -178,9 +178,9 @@ class TD3(BaseRLModel):
self.train(episode_timesteps, batch_size=self.batch_size, policy_freq=self.policy_freq)
# Evaluate episode
if 0 < eval_freq <= timesteps_since_eval:
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
timesteps_since_eval %= eval_freq
mean_reward, _ = evaluate_policy(self, self.env, n_eval_episodes)
mean_reward, _ = evaluate_policy(self, eval_env, n_eval_episodes)
evaluations.append(mean_reward)
if self.verbose > 0:
print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1]))