From 8acac6b0f4f7aca7b413af95c7e0151792cf4fae Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Mon, 3 Feb 2020 18:31:13 +0100 Subject: [PATCH] Update docstring --- torchy_baselines/cem_rl/cem_rl.py | 2 +- torchy_baselines/common/base_class.py | 41 ++++++++++++++++----------- torchy_baselines/sac/sac.py | 2 +- torchy_baselines/td3/td3.py | 2 +- 4 files changed, 27 insertions(+), 20 deletions(-) diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 653b3e2..9a2defa 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -164,7 +164,7 @@ class CEMRL(TD3): rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, n_steps=-1, action_noise=self.action_noise, - deterministic=False, callback=callback, + callback=callback, learning_starts=self.learning_starts, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num, diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index e46a462..5cb75d7 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -699,6 +699,12 @@ class OffPolicyRLModel(BaseRLModel): self.ep_info_buffer = None # type: deque self.use_sde_at_warmup = use_sde_at_warmup + def save_replay_buffer(self): + pass + + def load_replay_buffer(self, path): + pass + def collect_rollouts(self, env: VecEnv, # Type hint as string to avoid circular import @@ -706,7 +712,6 @@ class OffPolicyRLModel(BaseRLModel): n_episodes: int = 1, n_steps: int = -1, action_noise: Optional[ActionNoise] = None, - deterministic: bool = False, learning_starts: int = 0, replay_buffer: Optional[ReplayBuffer] = None, obs: Optional[np.ndarray] = None, @@ -715,23 +720,27 @@ class OffPolicyRLModel(BaseRLModel): """ Collect rollout using the current policy (and possibly fill the replay buffer) - :param env: (VecEnv) - :param n_episodes: (int) - :param n_steps: (int) - :param action_noise: (ActionNoise) - :param deterministic: (bool) - :param callback: (BaseCallback) - :param learning_starts: (int) + :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 obs: (np.ndarray) - :param episode_num: (int) - :param log_interval: (int) + :param obs: (np.ndarray) Last observation from the environment + :param episode_num: (int) Episode index + :param log_interval: (int) Log data every `log_interval` episodes """ - episode_rewards = [] - total_timesteps = [] + episode_rewards, total_timesteps = [], [] total_steps, total_episodes = 0, 0 - assert isinstance(env, VecEnv) - assert env.num_envs == 1 + + assert isinstance(env, VecEnv), "You must pass a VecEnv" + assert env.num_envs == 1, "OffPolicyRLModel only support single environment" # Retrieve unnormalized observation for saving into the buffer if self._vec_normalize_env is not None: @@ -749,8 +758,6 @@ class OffPolicyRLModel(BaseRLModel): while total_steps < n_steps or total_episodes < n_episodes: done = False - # Reset environment: not needed for VecEnv - # obs = env.reset() episode_reward, episode_timesteps = 0.0, 0 while not done: diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index 4d86cde..9617ec7 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -272,7 +272,7 @@ class SAC(OffPolicyRLModel): while self.num_timesteps < total_timesteps: rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, n_steps=self.train_freq, action_noise=self.action_noise, - deterministic=False, callback=callback, + callback=callback, learning_starts=self.learning_starts, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num, diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index a52a751..c485b4d 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -267,7 +267,7 @@ class TD3(OffPolicyRLModel): rollout = self.collect_rollouts(self.env, n_episodes=self.n_episodes_rollout, n_steps=self.train_freq, action_noise=self.action_noise, - deterministic=False, callback=callback, + callback=callback, learning_starts=self.learning_starts, replay_buffer=self.replay_buffer, obs=obs, episode_num=episode_num,