def collect_rollouts(
self,
env: VecEnv,
callback: BaseCallback,
train_freq: TrainFreq,
replay_buffer: ReplayBuffer,
action_noise: Optional[ActionNoise] = None,
learning_starts: int = 0,
log_interval: Optional[int] = None,
) -> RolloutReturn:
"""
Collect experiences and store them into a ``ReplayBuffer``.
:param env: The training environment
:param callback: Callback that will be called at each step
(and at the beginning and end of the rollout)
:param train_freq: How much experience to collect
by doing rollouts of current policy.
Either ``TrainFreq(<n>, TrainFrequencyUnit.STEP)``
or ``TrainFreq(<n>, TrainFrequencyUnit.EPISODE)``
with ``<n>`` being an integer greater than 0.
:param action_noise: Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC.
:param learning_starts: Number of steps before learning for the warm-up phase.
:param replay_buffer:
:param log_interval: Log data every ``log_interval`` episodes
:return:
"""
episode_rewards, total_timesteps = [], []
num_collected_steps, num_collected_episodes = 0, 0
assert isinstance(env, VecEnv), "You must pass a VecEnv"
assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
assert train_freq.frequency > 0, "Should at least collect one step or episode."
if self.use_sde:
self.actor.reset_noise()
callback.on_rollout_start()
continue_training = True
while should_collect_more_steps(train_freq, num_collected_steps, num_collected_episodes):
done = False
episode_reward, episode_timesteps = 0.0, 0
while not done:
if self.use_sde and self.sde_sample_freq > 0 and num_collected_steps % self.sde_sample_freq == 0:
# Sample a new noise matrix
self.actor.reset_noise()
# Select action randomly or according to policy
action, buffer_action = self._sample_action(learning_starts, action_noise)
# Rescale and perform action
new_obs, reward, done, infos = env.step(action)
self.num_timesteps += 1
episode_timesteps += 1
num_collected_steps += 1
# Give access to local variables
callback.update_locals(locals())
# Only stop training if return value is False, not when it is None.
if callback.on_step() is False:
return RolloutReturn(0.0, num_collected_steps, num_collected_episodes, continue_training=False)
episode_reward += reward
# Retrieve reward and episode length if using Monitor wrapper
self._update_info_buffer(infos, done)
# Store data in replay buffer (normalized action and unnormalized observation)
self._store_transition(replay_buffer, buffer_action, new_obs, reward, done, infos)
self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps)
# For DQN, check if the target network should be updated
# and update the exploration schedule
# For SAC/TD3, the update is done as the same time as the gradient update
# see https://github.com/hill-a/stable-baselines/issues/900
self._on_step()
if not should_collect_more_steps(train_freq, num_collected_steps, num_collected_episodes):
break
if done:
num_collected_episodes += 1
self._episode_num += 1
episode_rewards.append(episode_reward)
total_timesteps.append(episode_timesteps)
if action_noise is not None:
action_noise.reset()
# Log training infos
if log_interval is not None and self._episode_num % log_interval == 0:
self._dump_logs()
mean_reward = np.mean(episode_rewards) if num_collected_episodes > 0 else 0.0
callback.on_rollout_end()
return RolloutReturn(mean_reward, num_collected_steps, num_collected_episodes, continue_training)
def collect_rollouts(self,
env: VecEnv,
# Type hint as string to avoid circular import
callback: 'BaseCallback',
n_episodes: int = 1,
n_steps: int = -1,
action_noise: Optional[ActionNoise] = None,
learning_starts: int = 0,
replay_buffer: Optional[ReplayBuffer] = None,
log_interval: Optional[int] = None) -> RolloutReturn:
"""
Collect rollout using the current policy (and possibly fill the replay buffer)
:param env: (VecEnv) The training environment
:param n_episodes: (int) Number of episodes to use to collect rollout data
You can also specify a ``n_steps`` instead
:param n_steps: (int) Number of steps to use to collect rollout data
You can also specify a ``n_episodes`` instead.
:param action_noise: (Optional[ActionNoise]) Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC.
:param callback: (BaseCallback) Callback that will be called at each step
(and at the beginning and end of the rollout)
:param learning_starts: (int) Number of steps before learning for the warm-up phase.
:param replay_buffer: (ReplayBuffer)
:param log_interval: (int) Log data every ``log_interval`` episodes
:return: (RolloutReturn)
"""
episode_rewards, total_timesteps = [], []
total_steps, total_episodes = 0, 0
assert isinstance(env, VecEnv), "You must pass a VecEnv"
assert env.num_envs == 1, "OffPolicyRLModel only support single environment"
if self.use_sde:
self.actor.reset_noise()
callback.on_rollout_start()
continue_training = True
while total_steps < n_steps or total_episodes < n_episodes:
done = False
episode_reward, episode_timesteps = 0.0, 0
while not done:
if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
# Sample a new noise matrix
self.actor.reset_noise()
# Select action randomly or according to policy
if self.num_timesteps < learning_starts and not (self.use_sde and self.use_sde_at_warmup):
# Warmup phase
unscaled_action = np.array([self.action_space.sample()])
else:
# Note: we assume that the policy uses tanh to scale the action
# We use non-deterministic action in the case of SAC, for TD3, it does not matter
unscaled_action, _ = self.predict(self._last_obs, deterministic=False)
# Rescale the action from [low, high] to [-1, 1]
if isinstance(self.action_space, gym.spaces.Box):
scaled_action = self.policy.scale_action(unscaled_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
scaled_action = np.clip(scaled_action + action_noise(), -1, 1)
# We store the scaled action in the buffer
buffer_action = scaled_action
action = self.policy.unscale_action(scaled_action)
else:
# Discrete case, no need to normalize or clip
buffer_action = unscaled_action
action = buffer_action
# Rescale and perform action
new_obs, reward, done, infos = env.step(action)
# Only stop training if return value is False, not when it is None.
if callback.on_step() is False:
return RolloutReturn(0.0, total_steps, total_episodes, continue_training=False)
episode_reward += reward
# Retrieve reward and episode length if using Monitor wrapper
self._update_info_buffer(infos, done)
# Store data in replay buffer
if replay_buffer is not None:
# 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
self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward
replay_buffer.add(self._last_original_obs, new_obs_, buffer_action, reward_, done)
self._last_obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
self._last_original_obs = new_obs_
self.num_timesteps += 1
episode_timesteps += 1
total_steps += 1
if 0 < n_steps <= total_steps:
break
if done:
total_episodes += 1
self._episode_num += 1
episode_rewards.append(episode_reward)
total_timesteps.append(episode_timesteps)
if action_noise is not None:
action_noise.reset()
# Display training infos
if self.verbose >= 1 and log_interval is not None and self._episode_num % log_interval == 0:
fps = int(self.num_timesteps / (time.time() - self.start_time))
self.logger.logkv("episodes", self._episode_num)
if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:
self.logger.logkv('ep_rew_mean', self.safe_mean([ep_info['r'] for ep_info in self.ep_info_buffer]))
self.logger.logkv('ep_len_mean', self.safe_mean([ep_info['l'] for ep_info in self.ep_info_buffer]))
self.logger.logkv("fps", fps)
self.logger.logkv('time_elapsed', int(time.time() - self.start_time))
self.logger.logkv("total timesteps", self.num_timesteps)
if self.use_sde:
self.logger.logkv("std", (self.actor.get_std()).mean().item())
if len(self.ep_success_buffer) > 0:
self.logger.logkv('success rate', self.safe_mean(self.ep_success_buffer))
self.logger.dumpkvs()
mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0
callback.on_rollout_end()
return RolloutReturn(mean_reward, total_steps, total_episodes, continue_training)
def collect_rollouts(self,
env: VecEnv,
callback: BaseCallback,
n_episodes: int = 1,
n_steps: int = -1,
action_noise: Optional[ActionNoise] = None,
learning_starts: int = 0,
replay_buffer: Optional[ReplayBuffer] = None,
log_interval: Optional[int] = None) -> RolloutReturn:
"""
Collect experiences and store them into a ReplayBuffer.
:param env: (VecEnv) The training environment
:param callback: (BaseCallback) Callback that will be called at each step
(and at the beginning and end of the rollout)
:param n_episodes: (int) Number of episodes to use to collect rollout data
You can also specify a ``n_steps`` instead
:param n_steps: (int) Number of steps to use to collect rollout data
You can also specify a ``n_episodes`` instead.
:param action_noise: (Optional[ActionNoise]) Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC.
:param learning_starts: (int) Number of steps before learning for the warm-up phase.
:param replay_buffer: (ReplayBuffer)
:param log_interval: (int) Log data every ``log_interval`` episodes
:return: (RolloutReturn)
"""
episode_rewards, total_timesteps = [], []
total_steps, total_episodes = 0, 0
assert isinstance(env, VecEnv), "You must pass a VecEnv"
assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
if self.use_sde:
self.actor.reset_noise()
callback.on_rollout_start()
continue_training = True
while total_steps < n_steps or total_episodes < n_episodes:
done = False
episode_reward, episode_timesteps = 0.0, 0
while not done:
if self.use_sde and self.sde_sample_freq > 0 and total_steps % self.sde_sample_freq == 0:
# Sample a new noise matrix
self.actor.reset_noise()
# Select action randomly or according to policy
action, buffer_action, is_random_action = self._sample_action(learning_starts, action_noise)
# print("action/buffer action/israndomaction ", action, buffer_action,is_random_action)
# action, buffer_action = self._sample_action(learning_starts, action_noise)
# Rescale and perform action
new_obs, reward, done, infos = env.step(action)
# print("Reward: ", reward)
# print("Observation:", new_obs)
self.period_counter += 1
self.current_observation = new_obs
# Only stop training if return value is False, not when it is None.
if callback.on_step() is False:
return RolloutReturn(0.0, total_steps, total_episodes, continue_training=False)
episode_reward += reward
# Retrieve reward and episode length if using Monitor wrapper
self._update_info_buffer(infos, done)
# Store data in replay buffer
if replay_buffer is not None:
# 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
self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward
replay_buffer.add(self._last_original_obs, new_obs_, buffer_action, reward_,
done)
old_observation = self._last_original_obs
self._last_obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
self._last_original_obs = new_obs_
self.num_timesteps += 1
episode_timesteps += 1
total_steps += 1
self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps)
# For DQN, check if the target network should be updated
# and update the exploration schedule
# For SAC/TD3, the update is done as the same time as the gradient update
# see https://github.com/hill-a/stable-baselines/issues/900
self._on_step()
if 0 < n_steps <= total_steps:
break
# Compute the target Q values
target_st1 = self.q_net_target(th.tensor(new_obs))
# Follow greedy policy: use the one with the highest value
target_st1, _ = target_st1.max(dim=1)
# Avoid potential broadcast issue
target_st1 = float(target_st1.reshape(-1, 1))
# Compute the target Q values
target_st = self.q_net_target(th.tensor(old_observation))
# Follow greedy policy: use the one with the highest value
target_st, _ = target_st.max(dim=1)
# Avoid potential broadcast issue -> IS THIS NECESSARY?
target_st = float(target_st.reshape(-1, 1))
# TODO: Welchen der drei indices von target_st nehmen wir? ->maximalen
# TODO 2: laut paper wird zuerst rho berechnet, aber hier machen wir es einen step verzögert
# buffer_action.astype(int)[0]
if is_random_action == 0:
decayed_alpha = self.exp_decay_alpha()
self.rho = (1 - decayed_alpha) * self.rho + decayed_alpha * (reward_ + target_st1 - target_st)
# Fixed observation for debugging purposes
# TODO: GENERATE RANDOM NUMBERS OR COME UP WITH BETTER NUMBERS
obs2 = th.tensor([[10., 0., 0., 0., 5., 0., 7., 1., 0., 0., 5., 1., 6., 1.,
0., 0., 3., 0., 8., 0., 0., 0., 3., 2., 14., 2., 0., 0.,
4., 2., 17., 1., 0., 0., 3., 1., 10., 0., 0., 0., 5., 0., 7., 1., 0., 0., 5., 1., 6.,
1.,
0., 0., 3., 0., 8., 0., 0., 0., 3., 2., 14., 2., 0., 0.,
4., 2., 17., 1., 0., 0., 3., 1.,
10., 0., 0., 0., 5., 0., 7., 1., 0., 0., 5., 1., 6., 1.,
0., 0., 3., 0., 8., 0., 0., 0., 3., 2.,
10., 0., 0., 0., 5., 0., 7., 1., 0., 0., 5., 1.,
10., 0., 0., 0., 5., 0., 7., 1., 0., 0., 5., 1.,
10., 0., 0., 0., 5., 0., 7., 1., 0., 0., 5., 1.
]])
fix_observation = self.q_net._predict(obs2)[1][0]
with open('../' + 'q_values_learned_results.csv', mode='a') as results_CSV:
results_writer = csv.writer(results_CSV, delimiter='\t', quotechar='"', quoting=csv.QUOTE_MINIMAL)
results_writer.writerow(
[self.period_counter, float(fix_observation[0]), float(fix_observation[1]),
float(fix_observation[2])])
# Write reward to CSV file after each period
with open('../' + 'rewards_per_period.csv', mode='a') as rewards_per_period_CSV:
results_writer = csv.writer(rewards_per_period_CSV, delimiter='\t', quotechar='"',
quoting=csv.QUOTE_MINIMAL)
results_writer.writerow([self.period_counter, float(reward_), float(self.rho)])
if done:
total_episodes += 1
self._episode_num += 1
episode_rewards.append(episode_reward)
total_timesteps.append(episode_timesteps)
if action_noise is not None:
action_noise.reset()
# Log training infos
if log_interval is not None and self._episode_num % log_interval == 0:
self._dump_logs()
mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0
callback.on_rollout_end()
return RolloutReturn(mean_reward, total_steps, total_episodes, continue_training)
def collect_rollouts(
self,
env: VecEnv,
callback: BaseCallback,
n_episodes: int = 1,
n_steps: int = -1,
action_noise: Optional[ActionNoise] = None,
learning_starts: int = 0,
replay_buffer: Optional[ReplayBuffer] = None,
log_interval: Optional[int] = None,
) -> RolloutReturn:
"""
Collect experiences and store them into a ``ReplayBuffer``.
:param env: The training environment
:param callback: Callback that will be called at each step
(and at the beginning and end of the rollout)
:param n_episodes: Number of episodes to use to collect rollout data
You can also specify a ``n_steps`` instead
:param n_steps: Number of steps to use to collect rollout data
You can also specify a ``n_episodes`` instead.
:param action_noise: Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC.
:param learning_starts: Number of steps before learning for the warm-up phase.
:param replay_buffer:
:param log_interval: Log data every ``log_interval`` episodes
:return:
"""
episode_rewards, total_timesteps = [], []
total_steps, total_episodes = 0, 0
assert isinstance(env, VecEnv), "You must pass a VecEnv"
assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
if self.use_sde:
self.actor.reset_noise()
callback.on_rollout_start()
continue_training = True
while total_steps < n_steps or total_episodes < n_episodes:
done = False
episode_reward, episode_timesteps = 0.0, 0
while not done:
if self.use_sde and self.sde_sample_freq > 0 and total_steps % self.sde_sample_freq == 0:
# Sample a new noise matrix
self.actor.reset_noise()
# Select action randomly or according to policy
action, buffer_action = self._sample_action(
learning_starts, action_noise)
# Rescale and perform action
new_obs, reward, done, infos = env.step(action)
self.num_timesteps += 1
episode_timesteps += 1
total_steps += 1
# Give access to local variables
callback.update_locals(locals())
# Only stop training if return value is False, not when it is None.
if callback.on_step() is False:
return RolloutReturn(0.0,
total_steps,
total_episodes,
continue_training=False)
episode_reward += reward
# Retrieve reward and episode length if using Monitor wrapper
self._update_info_buffer(infos, done)
# Store data in replay buffer
if replay_buffer is not None:
# 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
self._last_original_obs, new_obs_, reward_ = self._last_obs, new_obs, reward
replay_buffer.add(self._last_original_obs, new_obs_,
buffer_action, reward_, done)
self._last_obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
self._last_original_obs = new_obs_
self._update_current_progress_remaining(
self.num_timesteps, self._total_timesteps)
# For DQN, check if the target network should be updated
# and update the exploration schedule
# For SAC/TD3, the update is done as the same time as the gradient update
# see https://github.com/hill-a/stable-baselines/issues/900
self._on_step()
if 0 < n_steps <= total_steps:
break
if done:
total_episodes += 1
self._episode_num += 1
episode_rewards.append(episode_reward)
total_timesteps.append(episode_timesteps)
if action_noise is not None:
action_noise.reset()
# Log training infos
if log_interval is not None and self._episode_num % log_interval == 0:
self._dump_logs()
mean_reward = np.mean(episode_rewards) if total_episodes > 0 else 0.0
callback.on_rollout_end()
return RolloutReturn(mean_reward, total_steps, total_episodes,
continue_training)
def collect_rollouts(
self,
env: VecEnv,
callback: BaseCallback,
train_freq: TrainFreq,
action_noise: Optional[ActionNoise] = None,
learning_starts: int = 0,
log_interval: Optional[int] = None,
) -> RolloutReturn:
"""
Collect experiences and store them into a ReplayBuffer.
:param env: The training environment
:param callback: Callback that will be called at each step
(and at the beginning and end of the rollout)
:param train_freq: How much experience to collect
by doing rollouts of current policy.
Either ``TrainFreq(<n>, TrainFrequencyUnit.STEP)``
or ``TrainFreq(<n>, TrainFrequencyUnit.EPISODE)``
with ``<n>`` being an integer greater than 0.
:param action_noise: Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC.
:param learning_starts: Number of steps before learning for the warm-up phase.
:param log_interval: Log data every ``log_interval`` episodes
:return:
"""
episode_rewards, total_timesteps = [], []
num_collected_steps, num_collected_episodes = 0, 0
assert isinstance(env, VecEnv), "You must pass a VecEnv"
assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
assert train_freq.frequency > 0, "Should at least collect one step or episode."
if self.model.use_sde:
self.actor.reset_noise()
callback.on_rollout_start()
continue_training = True
while should_collect_more_steps(train_freq, num_collected_steps,
num_collected_episodes):
done = False
episode_reward, episode_timesteps = 0.0, 0
while not done:
# concatenate observation and (desired) goal
observation = self._last_obs
self._last_obs = ObsDictWrapper.convert_dict(observation)
if (self.model.use_sde and self.model.sde_sample_freq > 0 and
num_collected_steps % self.model.sde_sample_freq == 0):
# Sample a new noise matrix
self.actor.reset_noise()
# Select action randomly or according to policy
self.model._last_obs = self._last_obs
action, buffer_action = self._sample_action(
learning_starts, action_noise)
# Perform action
new_obs, reward, done, infos = env.step(action)
self.num_timesteps += 1
self.model.num_timesteps = self.num_timesteps
episode_timesteps += 1
num_collected_steps += 1
# Only stop training if return value is False, not when it is None.
if callback.on_step() is False:
return RolloutReturn(0.0,
num_collected_steps,
num_collected_episodes,
continue_training=False)
episode_reward += reward
# Retrieve reward and episode length if using Monitor wrapper
self._update_info_buffer(infos, done)
self.model.ep_info_buffer = self.ep_info_buffer
self.model.ep_success_buffer = self.ep_success_buffer
# == Store transition in the replay buffer and/or in the episode storage ==
if self._vec_normalize_env is not None:
# Store only the unnormalized version
new_obs_ = self._vec_normalize_env.get_original_obs()
reward_ = self._vec_normalize_env.get_original_reward()
else:
# Avoid changing the original ones
self._last_original_obs, new_obs_, reward_ = observation, new_obs, reward
self.model._last_original_obs = self._last_original_obs
# As the VecEnv resets automatically, new_obs is already the
# first observation of the next episode
if done and infos[0].get("terminal_observation") is not None:
next_obs = infos[0]["terminal_observation"]
# VecNormalize normalizes the terminal observation
if self._vec_normalize_env is not None:
next_obs = self._vec_normalize_env.unnormalize_obs(
next_obs)
else:
next_obs = new_obs_
if self.online_sampling:
self.replay_buffer.add(self._last_original_obs, next_obs,
buffer_action, reward_, done, infos)
else:
# concatenate observation with (desired) goal
flattened_obs = ObsDictWrapper.convert_dict(
self._last_original_obs)
flattened_next_obs = ObsDictWrapper.convert_dict(next_obs)
# add to replay buffer
self.replay_buffer.add(flattened_obs, flattened_next_obs,
buffer_action, reward_, done)
# add current transition to episode storage
self._episode_storage.add(self._last_original_obs,
next_obs, buffer_action, reward_,
done, infos)
self._last_obs = new_obs
self.model._last_obs = self._last_obs
# Save the unnormalized new observation
if self._vec_normalize_env is not None:
self._last_original_obs = new_obs_
self.model._last_original_obs = self._last_original_obs
self.model._update_current_progress_remaining(
self.num_timesteps, self._total_timesteps)
# For DQN, check if the target network should be updated
# and update the exploration schedule
# For SAC/TD3, the update is done as the same time as the gradient update
# see https://github.com/hill-a/stable-baselines/issues/900
self.model._on_step()
self.episode_steps += 1
if not should_collect_more_steps(train_freq,
num_collected_steps,
num_collected_episodes):
break
if done or self.episode_steps >= self.max_episode_length:
if self.online_sampling:
self.replay_buffer.store_episode()
else:
self._episode_storage.store_episode()
# sample virtual transitions and store them in replay buffer
self._sample_her_transitions()
# clear storage for current episode
self._episode_storage.reset()
num_collected_episodes += 1
self._episode_num += 1
self.model._episode_num = self._episode_num
episode_rewards.append(episode_reward)
total_timesteps.append(episode_timesteps)
if action_noise is not None:
action_noise.reset()
# Log training infos
if log_interval is not None and self._episode_num % log_interval == 0:
self._dump_logs()
self.episode_steps = 0
mean_reward = np.mean(
episode_rewards) if num_collected_episodes > 0 else 0.0
callback.on_rollout_end()
return RolloutReturn(mean_reward, num_collected_steps,
num_collected_episodes, continue_training)
def collect_rollouts(
self,
env: VecEnv,
callback: BaseCallback,
train_freq: TrainFreq,
buffer: TrajectoryBuffer,
action_noise: Optional[ActionNoise] = None,
learning_starts: int = 0,
log_interval: Optional[int] = None,
) -> RolloutReturn:
"""
Collect experiences and store them into a ``TrajectoryBuffer``.
:param env: The training environment
:param callback: Callback that will be called at each step
(and at the beginning and end of the rollout)
:param train_freq: How much experience to collect
by doing rollouts of current policy.
Either ``TrainFreq(<n>, TrainFrequencyUnit.STEP)``
or ``TrainFreq(<n>, TrainFrequencyUnit.EPISODE)``
with ``<n>`` being an integer greater than 0.
:param action_noise: Action noise that will be used for exploration
Required for deterministic policy (e.g. TD3). This can also be used
in addition to the stochastic policy for SAC.
:param learning_starts: Number of steps before learning for the warm-up phase.
:param trajectory_buffer:
:param log_interval: Log data every ``log_interval`` episodes
:return:
"""
episode_rewards, total_timesteps = [], []
num_collected_steps, num_collected_episodes = 0, 0
assert isinstance(env, VecEnv), "You must pass a VecEnv"
# assert env.num_envs == 1, "OffPolicyAlgorithm only support single environment"
assert train_freq.frequency > 0, "Should at least collect one step or episode."
if self.use_sde:
self.actor.reset_noise()
callback.on_rollout_start()
continue_training = True
self.rollout_buffer.reset()
done = np.array([False for i in range(self.n_envs)])
episode_reward, episode_timesteps = [0.0 for i in range(self.n_envs)], [0 for i in range(self.n_envs)]
if train_freq.unit == TrainFrequencyUnit.STEP:
self.trajectories = [Trajectory(self.device) for i in range(self.n_envs)]
while True:
ms = [0]
get_ms(ms)
if self.use_sde and self.sde_sample_freq > 0 and num_collected_steps % self.sde_sample_freq == 0:
# Sample a new noise matrix
self.actor.reset_noise()
# Select action randomly or according to policy
with th.no_grad():
# action, buffer_action = self._sample_action(learning_starts, action_noise, use_behav=False)
# log_probs = self.policy.get_action_log_probs(th.tensor(np.array(self._last_obs)).to(self.device), th.tensor(np.array([action])).T.to(self.device), use_behav=False)
action, buffer_action = self._sample_action(learning_starts, action_noise, use_behav=True)
log_probs = self.policy.get_action_log_probs(th.tensor(np.array(self._last_obs)).to(self.device), th.tensor(np.array([action])).T.to(self.device), use_behav=True)
prob = th.exp(log_probs)
prob = (1 - self.exploration_rate) * prob + (self.exploration_rate) * (1.0 / self.action_space.n)
prob = prob.cpu().numpy()
if (prob > 1).any():
print("prob > 1!!! => Code in offpac.py")
print(prob)
print(th.tensor(log_probs))
exit()
new_obs, reward, done, infos = env.step(action)
with th.no_grad():
if self.use_v_net:
latent_pi, latent_vf, latent_sde = self.policy._get_latent(th.tensor(self._last_obs).to(self.device))
values = self.value_net(latent_vf).detach()
else:
values = self.policy.compute_value(th.tensor(self._last_obs).to(self.device), use_target_v=False).detach()
# self.rollout_buffer.add(self._last_obs, action.reshape(-1, 1), reward, self._last_episode_starts, values, log_probs.flatten())
self.num_timesteps += env.num_envs
num_collected_steps += env.num_envs
# Give access to local variables
callback.update_locals(locals())
# Only stop training if return value is False, not when it is None.
if callback.on_step() is False:
return RolloutReturn(0.0, num_collected_steps, num_collected_episodes, continue_training=False)
episode_reward += reward
# Retrieve reward and episode length if using Monitor wrapper
self._update_info_buffer(infos, done)
for i in range(len(self.trajectories)):
# trajectories[i].add(Transition(self._last_obs[i], action[i], reward[i], new_obs[i], done[i], prob[i]))
if done[i]:
if infos[i]['terminal_observation'].dtype == np.float64:
self.trajectories[i].add(Transition(self._last_obs[i], action[i], reward[i], infos[i]['terminal_observation'].astype(np.float32), done[i], prob[i]))
else:
self.trajectories[i].add(Transition(self._last_obs[i], action[i], reward[i], infos[i]['terminal_observation'], done[i], prob[i]))
else:
self.trajectories[i].add(Transition(self._last_obs[i], action[i], reward[i], new_obs[i], done[i], prob[i]))
self._last_obs = new_obs
self._last_episode_starts = done
self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps)
# For DQN, check if the target network should be updated
# and update the exploration schedule
# For SAC/TD3, the update is done as the same time as the gradient update
# see https://github.com/hill-a/stable-baselines/issues/900
self._on_step()
'''
if not should_collect_more_steps(train_freq, num_collected_steps, num_collected_episodes):
# even if the episdoe is not finished, we store the trajectory because no more steps can be performed
for traj_i, traj in enumerate(trajectories):
self._store_transition(buffer, traj)
total_timesteps.append(len(traj))
trajectories[traj_i] = Trajectory(self.device)
episode_rewards.append(episode_reward[traj_i])
episode_reward[traj_i] = 0.0
break
'''
# store transition of finished episode, but if not more steps can be collected, treat any trajectory as an episode
if done.any():
num_collected_episodes += np.sum(done)
self._episode_num += np.sum(done)
if log_interval is not None and self._episode_num % log_interval == 0:
self._dump_logs()
if train_freq.unit == TrainFrequencyUnit.STEP:
ending = not should_collect_more_steps(train_freq, num_collected_steps//self.n_envs, num_collected_episodes//self.n_envs)
# if ending, save all trajectories, otherwise only save done episode
if ending:
for traj_i, traj in enumerate(self.trajectories):
self._store_transition(buffer, traj)
# total_timesteps.append(len(traj)) # is this line affecting anything????
self.trajectories[traj_i] = Trajectory(self.device)
episode_rewards.append(episode_reward[traj_i])
episode_reward[traj_i] = 0.0
break
else:
if done.any():
traj_indexes = [i for i in np.arange(len(self.trajectories))[done]]
for traj_i in traj_indexes:
self._store_transition(buffer, self.trajectories[traj_i])
# total_timesteps.append(len(traj)) # is this line affecting anything????
self.trajectories[traj_i] = Trajectory(self.device)
episode_rewards.append(episode_reward[traj_i])
episode_reward[traj_i] = 0.0
elif train_freq.unit == TrainFrequencyUnit.EPISODE:
ending = not should_collect_more_steps(train_freq, num_collected_steps//self.n_envs, num_collected_episodes//self.n_envs)
if done.any():
# if ending, save all trajectories even if not finished
# if not ending:
traj_indexes = [i for i in np.arange(len(self.trajectories))[done]]
for traj_i in traj_indexes:
self._store_transition(buffer, self.trajectories[traj_i])
# total_timesteps.append(len(traj)) # is this line affecting anything????
self.trajectories[traj_i] = Trajectory(self.device)
episode_rewards.append(episode_reward[traj_i])
episode_reward[traj_i] = 0.0
'''
else:
_trajectories = trajectories
for traj_i, traj in enumerate(_trajectories):
self._store_transition(buffer, traj)
total_timesteps.append(len(traj)) # is this line affecting anything????
self.trajectories[traj_i] = Trajectory(self.device)
episode_rewards.append(episode_reward[traj_i])
episode_reward[traj_i] = 0.0
'''
if ending:
break
else:
print(train_freq.unit)
raise Exception("Weird train_freq.unit...")
exit(-1)
if done.any():
if action_noise is not None:
action_noise.reset()
with th.no_grad():
obs_tensor = th.as_tensor(new_obs).squeeze(1).to(self.device)
if self.use_v_net:
latent_pi, latent_vf, latent_sde = self.policy._get_latent(obs_tensor)
values = self.value_net(latent_vf).detach()
else:
values = self.policy.compute_value(obs_tensor, use_target_v=False)
self.rollout_buffer.compute_returns_and_advantage(last_values=values, dones=done)
mean_reward = np.mean(episode_rewards) if num_collected_episodes > 0 else 0.0
callback.on_rollout_end()
return RolloutReturn(mean_reward, num_collected_steps, num_collected_episodes, continue_training)