From dbe8cfceb6376c6915c657f91f7d25b23b6a8176 Mon Sep 17 00:00:00 2001 From: Stelios Tymvios <52372765+PartiallyTyped@users.noreply.github.com> Date: Fri, 17 Jul 2020 16:53:28 +0300 Subject: [PATCH 1/2] Optimized polyak updates (#106) * quick polyak updates * changelog * typing * reverted autoformatting * rerverted autofmt * Update stable_baselines3/common/utils.py Co-authored-by: Antonin RAFFIN * parameter names in test * cleanup * Merge branch 'master' into polyak * Update changelog * Apply suggestions from code review Co-authored-by: Antonin RAFFIN * Update stable_baselines3/common/utils.py * Update utils.py Co-authored-by: Antonin RAFFIN --- docs/misc/changelog.rst | 3 ++- stable_baselines3/common/utils.py | 24 +++++++++++++++++++++++- stable_baselines3/dqn/dqn.py | 5 ++--- stable_baselines3/sac/sac.py | 4 ++-- stable_baselines3/td3/td3.py | 9 +++------ stable_baselines3/version.txt | 2 +- tests/test_utils.py | 15 +++++++++++++++ 7 files changed, 48 insertions(+), 14 deletions(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index f49fb45..722e9e7 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,7 +3,7 @@ Changelog ========== -Pre-Release 0.8.0a4 (WIP) +Pre-Release 0.8.0a5 (WIP) ------------------------------ Breaking Changes: @@ -40,6 +40,7 @@ Others: - Split the ``collect_rollout()`` method for off-policy algorithms - Added ``_on_step()`` for off-policy base class - Optimized replay buffer size by removing the need of ``next_observations`` numpy array +- Optimized polyak updates (1.5-1.95 speedup) through inplace operations (@PartiallyTyped) - Switch to ``black`` codestyle and added ``make format``, ``make check-codestyle`` and ``commit-checks`` - Ignored errors from newer pytype version - Added a check when using ``gSDE`` diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 9c00f0e..839bc59 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -2,7 +2,7 @@ import glob import os import random from collections import deque -from typing import Callable, Optional, Union +from typing import Callable, Iterable, Optional, Union import gym import numpy as np @@ -284,3 +284,25 @@ def safe_mean(arr: Union[np.ndarray, list, deque]) -> np.ndarray: :return: """ return np.nan if len(arr) == 0 else np.mean(arr) + + +def polyak_update(params: Iterable[th.nn.Parameter], target_params: Iterable[th.nn.Parameter], tau: float) -> None: + """ + Perform a Polyak average update on ``target_params`` using ``params``: + target parameters are slowly updated towards the main parameters. + ``tau``, the soft update coefficient controls the interpolation: + ``tau=1`` corresponds to copying the parameters to the target ones whereas nothing happens when ``tau=0``. + The Polyak update is done in place, with ``no_grad``, and therefore does not create intermediate tensors, + or a computation graph, reducing memory cost and improving performance. We scale the target params + by ``1-tau`` (in-place), add the new weights, scaled by ``tau`` and store the result of the sum in the target + params (in place). + See https://github.com/DLR-RM/stable-baselines3/issues/93 + + :param params: (Iterable[th.nn.Parameter]) parameters to use to update the target params + :param target_params: (Iterable[th.nn.Parameter]) parameters to update + :param tau: (float) the soft update coefficient ("Polyak update", between 0 and 1) + """ + with th.no_grad(): + for param, target_param in zip(params, target_params): + target_param.data.mul_(1 - tau) + th.add(target_param.data, param.data, alpha=tau, out=target_param.data) diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 55d4fb6..0b93c06 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -7,7 +7,7 @@ from torch.nn import functional as F from stable_baselines3.common import logger from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback -from stable_baselines3.common.utils import get_linear_fn +from stable_baselines3.common.utils import get_linear_fn, polyak_update from stable_baselines3.dqn.policies import DQNPolicy @@ -138,8 +138,7 @@ class DQN(OffPolicyAlgorithm): This method is called in ``collect_rollout()`` after each step in the environment. """ if self.num_timesteps % self.target_update_interval == 0: - for param, target_param in zip(self.q_net.parameters(), self.q_net_target.parameters()): - target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) + polyak_update(self.q_net.parameters(), self.q_net_target.parameters(), self.tau) self.exploration_rate = self.exploration_schedule(self._current_progress_remaining) logger.record("rollout/exploration rate", self.exploration_rate) diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 39b8e3a..7def773 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -8,6 +8,7 @@ from stable_baselines3.common import logger from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.utils import polyak_update from stable_baselines3.sac.policies import SACPolicy @@ -256,8 +257,7 @@ class SAC(OffPolicyAlgorithm): # Update target networks if gradient_step % self.target_update_interval == 0: - for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()): - target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) + polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau) self._n_updates += gradient_steps diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index 368bce1..8b7c902 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -7,6 +7,7 @@ from stable_baselines3.common import logger from stable_baselines3.common.noise import ActionNoise from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.utils import polyak_update from stable_baselines3.td3.policies import TD3Policy @@ -166,12 +167,8 @@ class TD3(OffPolicyAlgorithm): actor_loss.backward() self.actor.optimizer.step() - # Update the frozen target networks - for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()): - target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) - - for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()): - target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data) + polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau) + polyak_update(self.actor.parameters(), self.actor_target.parameters(), self.tau) self._n_updates += gradient_steps logger.record("train/n_updates", self._n_updates, exclude="tensorboard") diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index 9678ab2..87f68d6 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.8.0a4 +0.8.0a5 diff --git a/tests/test_utils.py b/tests/test_utils.py index 7f03c1d..4e1899f 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -4,6 +4,7 @@ import shutil import gym import numpy as np import pytest +import torch as th from stable_baselines3 import A2C from stable_baselines3.common.atari_wrappers import ClipRewardEnv @@ -11,6 +12,7 @@ from stable_baselines3.common.cmd_util import make_atari_env, make_vec_env from stable_baselines3.common.evaluation import evaluate_policy from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.noise import ActionNoise, OrnsteinUhlenbeckActionNoise, VectorizedActionNoise +from stable_baselines3.common.utils import polyak_update from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv @@ -152,3 +154,16 @@ def test_vec_noise(): vec.noises = [base] * (num_envs - 1) assert all(isinstance(noise, type(base)) for noise in vec.noises) assert len(vec.noises) == num_envs + + +def test_polyak(): + param1, param2 = th.nn.Parameter(th.ones((5, 5))), th.nn.Parameter(th.zeros((5, 5))) + target1, target2 = th.nn.Parameter(th.ones((5, 5))), th.nn.Parameter(th.zeros((5, 5))) + tau = 0.1 + polyak_update([param1], [param2], tau) + with th.no_grad(): + for param, target_param in zip([target1], [target2]): + target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data) + + assert th.allclose(param1, target1) + assert th.allclose(param2, target2) From 41b66e33ce850eb75f136a69e3373ed4fa97ee12 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Mon, 20 Jul 2020 11:20:12 +0200 Subject: [PATCH 2/2] [ci skip] Remove whitespace --- stable_baselines3/common/utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 839bc59..1c1b850 100644 --- a/stable_baselines3/common/utils.py +++ b/stable_baselines3/common/utils.py @@ -292,12 +292,12 @@ def polyak_update(params: Iterable[th.nn.Parameter], target_params: Iterable[th. target parameters are slowly updated towards the main parameters. ``tau``, the soft update coefficient controls the interpolation: ``tau=1`` corresponds to copying the parameters to the target ones whereas nothing happens when ``tau=0``. - The Polyak update is done in place, with ``no_grad``, and therefore does not create intermediate tensors, - or a computation graph, reducing memory cost and improving performance. We scale the target params - by ``1-tau`` (in-place), add the new weights, scaled by ``tau`` and store the result of the sum in the target - params (in place). + The Polyak update is done in place, with ``no_grad``, and therefore does not create intermediate tensors, + or a computation graph, reducing memory cost and improving performance. We scale the target params + by ``1-tau`` (in-place), add the new weights, scaled by ``tau`` and store the result of the sum in the target + params (in place). See https://github.com/DLR-RM/stable-baselines3/issues/93 - + :param params: (Iterable[th.nn.Parameter]) parameters to use to update the target params :param target_params: (Iterable[th.nn.Parameter]) parameters to update :param tau: (float) the soft update coefficient ("Polyak update", between 0 and 1)