diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 34da520..9841172 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -67,6 +67,7 @@ class BaseRLModel(object): 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 @@ -391,7 +392,7 @@ class BaseRLModel(object): obs_ = self._vec_normalize_env.get_original_obs() self.rollout_data = None - if hasattr(self, 'use_sde') and self.use_sde: + if self.use_sde: self.actor.reset_noise() # Reset rollout data self.rollout_data = {key: [] for key in ['observations', 'actions', 'rewards', 'dones']} @@ -405,23 +406,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: - if hasattr(self, 'use_sde'): - deterministic = not self.use_sde - 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 # 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 - 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 @@ -439,12 +446,12 @@ class BaseRLModel(object): # Avoid changing the original ones 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: # Assume only one env 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['dones'].append(np.array(done_bool[0]).copy()) @@ -480,7 +487,7 @@ class BaseRLModel(object): logger.logkv("fps", fps) logger.logkv('time_elapsed', int(time.time() - self.start_time)) 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.dumpkvs() diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 96fae45..7ef7fd7 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -22,10 +22,10 @@ class Actor(BaseNetwork): latent_pi = create_mlp(obs_dim, -1, net_arch, activation_fn, squash_out=False) self.latent_pi = nn.Sequential(*latent_pi) 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: # 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.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: 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(): latent_pi = self.latent_pi(obs) mean_actions = self.actor_net(latent_pi) @@ -68,8 +68,10 @@ class Actor(BaseNetwork): noise = th.mm(latent_pi.detach(), self.exploration_mat) if self.clip_noise is not None: noise = th.clamp(noise, -self.clip_noise, self.clip_noise) - # TODO: fix clipping with squashing ? - return th.clamp(self.actor_net(latent_pi) + noise, -1, 1) + # TODO: Replace with squashing -> need to account for that in the sde update + # 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: return self.actor_net(obs) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 2e312ac..5ecb38e 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -52,6 +52,7 @@ class TD3(BaseRLModel): Setting it to auto, the code will be run on the GPU if possible. :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance """ + def __init__(self, policy, env, buffer_size=int(1e6), learning_rate=1e-3, policy_delay=2, learning_starts=100, gamma=0.99, batch_size=100, 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) self.buffer_size = buffer_size - # TODO: accept callables self.learning_rate = learning_rate self.learning_starts = learning_starts self.train_freq = train_freq @@ -78,6 +78,7 @@ 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 @@ -204,10 +205,10 @@ class TD3(BaseRLModel): # self._update_learning_rate(self.policy.optimizer) # 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 - log_prob, entropy = self.actor.get_distribution_stats(obs, action) + log_prob, entropy = self.actor.evaluate_actions(obs, action) # Normalize returns # 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 self.train(gradient_steps, batch_size=self.batch_size, policy_delay=self.policy_delay) - # Evaluate episode if 0 < eval_freq <= timesteps_since_eval and eval_env is not None: timesteps_since_eval %= eval_freq