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

Fix clipped action when adapting noise with TD3

Browse source
This commit is contained in:
Antonin Raffin 2019-11-22 15:04:34 +01:00
parent 604a19fbc3
commit 03d2ab10f8
3 changed files with 29 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
torchy_baselines/common/base_class.py
View file

@ -67,6 +67,7 @@ class BaseRLModel(object):
self.action_noise = None self.action_noise = None
# Used for SDE only # Used for SDE only
self.rollout_data = None self.rollout_data = None
self.use_sde = False
# Track the training progress (from 1 to 0) # Track the training progress (from 1 to 0)
# this is used to update the learning rate # this is used to update the learning rate
self._current_progress = 1 self._current_progress = 1
@ -391,7 +392,7 @@ class BaseRLModel(object):
obs_ = self._vec_normalize_env.get_original_obs() obs_ = self._vec_normalize_env.get_original_obs()
self.rollout_data = None self.rollout_data = None
if hasattr(self, 'use_sde') and self.use_sde: if self.use_sde:
self.actor.reset_noise() self.actor.reset_noise()
# Reset rollout data # Reset rollout data
self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']} self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']}
@ -405,23 +406,29 @@ class BaseRLModel(object):
while not done: while not done:
# Select action randomly or according to policy # Select action randomly or according to policy
if num_timesteps < learning_starts: if num_timesteps < learning_starts:
action = np.array([self.action_space.sample()]) # Warmup phase
unscaled_action = np.array([self.action_space.sample()])
else: else:
if hasattr(self, 'use_sde'): unscaled_action = self.predict(obs, deterministic=not self.use_sde)
deterministic = not self.use_sde
action = self.predict(obs, deterministic=deterministic)
# Rescale the action from [low, high] to [-1, 1] # Rescale the action from [low, high] to [-1, 1]
action = self.scale_action(action) scaled_action = self.scale_action(unscaled_action)
if self.use_sde:
# When using SDE, the action can be out of bounds
# TODO: fix with squashing and account for that in the proba distribution
clipped_action = np.clip(scaled_action, -1, 1)
else:
clipped_action = scaled_action
# Add noise to the action (improve exploration) # Add noise to the action (improve exploration)
if action_noise is not None: if action_noise is not None:
# NOTE: in the original implementation of TD3, the noise was applied to the unscaled action # NOTE: in the original implementation of TD3, the noise was applied to the unscaled action
# Update(October 2019): Not anymore # Update(October 2019): Not anymore
action = np.clip(action + action_noise(), -1, 1) clipped_action = np.clip(clipped_action + action_noise(), -1, 1)
# Rescale and perform action # Rescale and perform action
new_obs, reward, done, infos = env.step(self.unscale_action(action)) new_obs, reward, done, infos = env.step(self.unscale_action(clipped_action))
done_bool = [float(done[0])] done_bool = [float(done[0])]
episode_reward += reward episode_reward += reward
@ -439,12 +446,12 @@ class BaseRLModel(object):
# Avoid changing the original ones # Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, reward obs_, new_obs_, reward_ = obs, new_obs, reward
replay_buffer.add(obs_, new_obs_, action, reward_, done_bool) replay_buffer.add(obs_, new_obs_, clipped_action, reward_, done_bool)
if self.rollout_data is not None: if self.rollout_data is not None:
# Assume only one env # Assume only one env
self.rollout_data['observations'].append(obs[0].copy()) self.rollout_data['observations'].append(obs[0].copy())
self.rollout_data['actions'].append(action[0].copy()) self.rollout_data['actions'].append(scaled_action[0].copy())
self.rollout_data['rewards'].append(reward[0].copy()) self.rollout_data['rewards'].append(reward[0].copy())
self.rollout_data['dones'].append(np.array(done_bool[0]).copy()) self.rollout_data['dones'].append(np.array(done_bool[0]).copy())
@ -480,7 +487,7 @@ class BaseRLModel(object):
logger.logkv("fps", fps) logger.logkv("fps", fps)
logger.logkv('time_elapsed', int(time.time() - self.start_time)) logger.logkv('time_elapsed', int(time.time() - self.start_time))
logger.logkv("total timesteps", num_timesteps) logger.logkv("total timesteps", num_timesteps)
if hasattr(self, 'use_sde') and self.use_sde: if self.use_sde:
logger.logkv("std", th.exp(self.actor.log_std).mean().item()) logger.logkv("std", th.exp(self.actor.log_std).mean().item())
logger.dumpkvs() logger.dumpkvs()

12
torchy_baselines/td3/policies.py
View file

@ -22,10 +22,10 @@ class Actor(BaseNetwork):
latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False) latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False)
self.latent_pi = nn.Sequential(*latent_pi) self.latent_pi = nn.Sequential(*latent_pi)
if full_std: if full_std:
self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init) self.log_std = nn.Parameter(th.ones(latent_dim, action_dim) * log_std_init, requires_grad=True)
else: else:
# Reduce the number of parameters: # Reduce the number of parameters:
self.log_std = nn.Parameter(th.ones(latent_dim, 1) * log_std_init) self.log_std = nn.Parameter(th.ones(latent_dim, 1) * log_std_init, requires_grad=True)
self.latent_dim = latent_dim self.latent_dim = latent_dim
self.actor_net = nn.Sequential(nn.Linear(net_arch[-1], action_dim), nn.Tanh()) self.actor_net = nn.Sequential(nn.Linear(net_arch[-1], action_dim), nn.Tanh())
@ -42,7 +42,7 @@ class Actor(BaseNetwork):
# Reduce the number of parameters: # Reduce the number of parameters:
return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std return th.ones((self.latent_dim, self.action_dim)).to(self.log_std.device) * self.log_std
def get_distribution_stats(self, obs, action): def evaluate_actions(self, obs, action):
with th.no_grad(): with th.no_grad():
latent_pi = self.latent_pi(obs) latent_pi = self.latent_pi(obs)
mean_actions = self.actor_net(latent_pi) mean_actions = self.actor_net(latent_pi)
@ -68,8 +68,10 @@ class Actor(BaseNetwork):
noise = th.mm(latent_pi.detach(), self.exploration_mat) noise = th.mm(latent_pi.detach(), self.exploration_mat)
if self.clip_noise is not None: if self.clip_noise is not None:
noise = th.clamp(noise, -self.clip_noise, self.clip_noise) noise = th.clamp(noise, -self.clip_noise, self.clip_noise)
# TODO: fix clipping with squashing ? # TODO: Replace with squashing -> need to account for that in the sde update
return th.clamp(self.actor_net(latent_pi) + noise, -1, 1) # return th.clamp(self.actor_net(latent_pi) + noise, -1, 1)
# NOTE: the clipping is done in the rollout for now
return self.actor_net(latent_pi) + noise
else: else:
return self.actor_net(obs) return self.actor_net(obs)

8
torchy_baselines/td3/td3.py
View file

@ -52,6 +52,7 @@ class TD3(BaseRLModel):
Setting it to auto, the code will be run on the GPU if possible. 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 :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance
""" """
def __init__(self, policy, env, buffer_size=int(1e6), learning_rate=1e-3, def __init__(self, policy, env, buffer_size=int(1e6), learning_rate=1e-3,
policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100,
train_freq=-1, gradient_steps=-1, n_episodes_rollout=1, train_freq=-1, gradient_steps=-1, n_episodes_rollout=1,
@ -64,7 +65,6 @@ class TD3(BaseRLModel):
create_eval_env=create_eval_env, seed=seed) create_eval_env=create_eval_env, seed=seed)
self.buffer_size = buffer_size self.buffer_size = buffer_size
# TODO: accept callables
self.learning_rate = learning_rate self.learning_rate = learning_rate
self.learning_starts = learning_starts self.learning_starts = learning_starts
self.train_freq = train_freq self.train_freq = train_freq
@ -78,6 +78,7 @@ class TD3(BaseRLModel):
self.target_noise_clip = target_noise_clip self.target_noise_clip = target_noise_clip
self.target_policy_noise = target_policy_noise self.target_policy_noise = target_policy_noise
# State Dependent Exploration
self.use_sde = use_sde self.use_sde = use_sde
self.sde_max_grad_norm = sde_max_grad_norm self.sde_max_grad_norm = sde_max_grad_norm
self.sde_ent_coef = sde_ent_coef self.sde_ent_coef = sde_ent_coef
@ -204,10 +205,10 @@ class TD3(BaseRLModel):
# self._update_learning_rate(self.policy.optimizer) # self._update_learning_rate(self.policy.optimizer)
# Unpack # Unpack
obs, action, returns = self.rollout_data['observations'], self.rollout_data['actions'], self.rollout_data['returns'] obs, action, returns = [self.rollout_data[key] for key in ['observations', 'actions', 'returns']]
# TODO: avoid second computation of everything because of the gradient # TODO: avoid second computation of everything because of the gradient
log_prob, entropy = self.actor.get_distribution_stats(obs, action) log_prob, entropy = self.actor.evaluate_actions(obs, action)
# Normalize returns # Normalize returns
# returns = (returns - returns.mean()) / (returns.std() + 1e-8) # returns = (returns - returns.mean()) / (returns.std() + 1e-8)
@ -281,7 +282,6 @@ class TD3(BaseRLModel):
gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps gradient_steps = self.gradient_steps if self.gradient_steps > 0 else episode_timesteps
self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay)
# Evaluate episode # Evaluate episode
if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: if 0 < eval_freq <= timesteps_since_eval and eval_env is not None:
timesteps_since_eval %= eval_freq timesteps_since_eval %= eval_freq