diff --git a/docs/guide/custom_policy.rst b/docs/guide/custom_policy.rst index 2ce9202..755baee 100644 --- a/docs/guide/custom_policy.rst +++ b/docs/guide/custom_policy.rst @@ -16,7 +16,7 @@ using ``policy_kwargs`` parameter: from stable_baselines3 import PPO - # Custom MLP policy of two layers of size 32 each with tanh activation function + # Custom MLP policy of two layers of size 32 each with Relu activation function policy_kwargs = dict(activation_fn=th.nn.ReLU, net_arch=[32, 32]) # Create the agent model = PPO("MlpPolicy", "CartPole-v1", policy_kwargs=policy_kwargs, verbose=1) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 0fa130a..f30ed95 100644 --- a/docs/misc/changelog.rst +++ b/docs/misc/changelog.rst @@ -3,9 +3,11 @@ Changelog ========== -Pre-Release 0.8.0a4 (WIP) +Pre-Release 0.8.0 (2020-08-03) ------------------------------ +**DQN, DDPG, bug fixes and performance matching for Atari games** + Breaking Changes: ^^^^^^^^^^^^^^^^^ - ``AtariWrapper`` and other Atari wrappers were updated to match SB2 ones @@ -31,6 +33,12 @@ Bug Fixes: ^^^^^^^^^^ - Fixed a bug in the ``close()`` method of ``SubprocVecEnv``, causing wrappers further down in the wrapper stack to not be closed. (@NeoExtended) - Fix target for updating q values in SAC: the entropy term was not conditioned by terminals states +- Use ``cloudpickle.load`` instead of ``pickle.load`` in ``CloudpickleWrapper``. (@shwang) +- Fixed a bug with orthogonal initialization when `bias=False` in custom policy (@rk37) +- Fixed approximate entropy calculation in PPO and A2C. (@andyshih12) +- Fixed DQN target network sharing feature extractor with the main network. +- Fixed storing correct ``dones`` in on-policy algorithm rollout collection. (@andyshih12) +- Fixed number of filters in final convolutional layer in NatureCNN to match original implementation. Deprecations: ^^^^^^^^^^^^^ @@ -41,10 +49,12 @@ 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`` - Removed codacy dependency from Dockerfile +- Added ``common.sb2_compat.RMSpropTFLike`` optimizer, which corresponds closer to the implementation of RMSprop from Tensorflow. Documentation: ^^^^^^^^^^^^^^ @@ -53,6 +63,7 @@ Documentation: - Added Unity reacher to the projects page (@koulakis) - Added PyBullet colab notebook - Fixed typo in PPO example code (@joeljosephjin) +- Fixed typo in custom policy doc (@RaphaelWag) @@ -355,4 +366,4 @@ And all the contributors: @Miffyli @dwiel @miguelrass @qxcv @jaberkow @eavelardev @ruifeng96150 @pedrohbtp @srivatsankrishnan @evilsocket @MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3 -@tirafesi @blurLake @koulakis @joeljosephjin +@tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag diff --git a/docs/modules/a2c.rst b/docs/modules/a2c.rst index 141fa0b..460d1a6 100644 --- a/docs/modules/a2c.rst +++ b/docs/modules/a2c.rst @@ -10,6 +10,14 @@ A synchronous, deterministic variant of `Asynchronous Advantage Actor Critic (A3 It uses multiple workers to avoid the use of a replay buffer. +.. warning:: + + If you find training unstable or want to match performance of stable-baselines A2C, consider using + ``RMSpropTFLike`` optimizer from ``stable_baselines3.common.sb2_compat.rmsprop_tf_like``. + You can change optimizer with ``A2C(policy_kwargs=dict(optimizer_class=RMSpropTFLike))``. + Read more `here `_. + + Notes ----- diff --git a/stable_baselines3/a2c/a2c.py b/stable_baselines3/a2c/a2c.py index c2c7b34..dc32476 100644 --- a/stable_baselines3/a2c/a2c.py +++ b/stable_baselines3/a2c/a2c.py @@ -116,6 +116,7 @@ class A2C(OnPolicyAlgorithm): # Update optimizer learning rate self._update_learning_rate(self.policy.optimizer) + # This will only loop once (get all data in one go) for rollout_data in self.rollout_buffer.get(batch_size=None): actions = rollout_data.actions @@ -141,7 +142,7 @@ class A2C(OnPolicyAlgorithm): # Entropy loss favor exploration if entropy is None: # Approximate entropy when no analytical form - entropy_loss = -log_prob.mean() + entropy_loss = -th.mean(-log_prob) else: entropy_loss = -th.mean(entropy) diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index 1ba7cc0..a3de8ce 100644 --- a/stable_baselines3/common/base_class.py +++ b/stable_baselines3/common/base_class.py @@ -123,6 +123,7 @@ class BaseAlgorithm(ABC): self.tensorboard_log = tensorboard_log self.lr_schedule = None # type: Optional[Callable] self._last_obs = None # type: Optional[np.ndarray] + self._last_dones = None # type: Optional[np.ndarray] # When using VecNormalize: self._last_original_obs = None # type: Optional[np.ndarray] self._episode_num = 0 @@ -474,6 +475,7 @@ class BaseAlgorithm(ABC): # Avoid resetting the environment when calling ``.learn()`` consecutive times if reset_num_timesteps or self._last_obs is None: self._last_obs = self.env.reset() + self._last_dones = np.zeros((self.env.num_envs,), dtype=np.bool) # Retrieve unnormalized observation for saving into the buffer if self._vec_normalize_env is not None: self._last_original_obs = self._vec_normalize_env.get_original_obs() diff --git a/stable_baselines3/common/on_policy_algorithm.py b/stable_baselines3/common/on_policy_algorithm.py index f84d18f..8c9cb8b 100644 --- a/stable_baselines3/common/on_policy_algorithm.py +++ b/stable_baselines3/common/on_policy_algorithm.py @@ -173,8 +173,9 @@ class OnPolicyAlgorithm(BaseAlgorithm): if isinstance(self.action_space, gym.spaces.Discrete): # Reshape in case of discrete action actions = actions.reshape(-1, 1) - rollout_buffer.add(self._last_obs, actions, rewards, dones, values, log_probs) + rollout_buffer.add(self._last_obs, actions, rewards, self._last_dones, values, log_probs) self._last_obs = new_obs + self._last_dones = dones rollout_buffer.compute_returns_and_advantage(values, dones=dones) diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index acadd88..01b788f 100644 --- a/stable_baselines3/common/policies.py +++ b/stable_baselines3/common/policies.py @@ -188,7 +188,8 @@ class BasePolicy(BaseModel): """ if isinstance(module, (nn.Linear, nn.Conv2d)): nn.init.orthogonal_(module.weight, gain=gain) - module.bias.data.fill_(0.0) + if module.bias is not None: + module.bias.data.fill_(0.0) @abstractmethod def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: diff --git a/stable_baselines3/common/sb2_compat/__init__.py b/stable_baselines3/common/sb2_compat/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py b/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py new file mode 100644 index 0000000..46ef4b0 --- /dev/null +++ b/stable_baselines3/common/sb2_compat/rmsprop_tf_like.py @@ -0,0 +1,126 @@ +import torch +from torch.optim import Optimizer + + +class RMSpropTFLike(Optimizer): + r"""Implements RMSprop algorithm with closer match to Tensorflow version. + + For reproducibility with original stable-baselines. Use this + version with e.g. A2C for stabler learning than with the PyTorch + RMSProp. Based on the PyTorch v1.5.0 implementation of RMSprop. + + See a more throughout conversion in pytorch-image-models repository: + https://github.com/rwightman/pytorch-image-models/blob/master/timm/optim/rmsprop_tf.py + + Changes to the original RMSprop: + - Move epsilon inside square root + - Initialize squared gradient to ones rather than zeros + + Proposed by G. Hinton in his + `course `_. + + The centered version first appears in `Generating Sequences + With Recurrent Neural Networks `_. + + The implementation here takes the square root of the gradient average before + adding epsilon (note that TensorFlow interchanges these two operations). The effective + learning rate is thus :math:`\alpha/(\sqrt{v} + \epsilon)` where :math:`\alpha` + is the scheduled learning rate and :math:`v` is the weighted moving average + of the squared gradient. + + Arguments: + params (iterable): iterable of parameters to optimize or dicts defining + parameter groups + lr (float, optional): learning rate (default: 1e-2) + momentum (float, optional): momentum factor (default: 0) + alpha (float, optional): smoothing constant (default: 0.99) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + centered (bool, optional) : if ``True``, compute the centered RMSProp, + the gradient is normalized by an estimation of its variance + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + + """ + + def __init__(self, params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False): + if not 0.0 <= lr: + raise ValueError("Invalid learning rate: {}".format(lr)) + if not 0.0 <= eps: + raise ValueError("Invalid epsilon value: {}".format(eps)) + if not 0.0 <= momentum: + raise ValueError("Invalid momentum value: {}".format(momentum)) + if not 0.0 <= weight_decay: + raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) + if not 0.0 <= alpha: + raise ValueError("Invalid alpha value: {}".format(alpha)) + + defaults = dict(lr=lr, momentum=momentum, alpha=alpha, eps=eps, centered=centered, weight_decay=weight_decay) + super(RMSpropTFLike, self).__init__(params, defaults) + + def __setstate__(self, state): + super(RMSpropTFLike, self).__setstate__(state) + for group in self.param_groups: + group.setdefault("momentum", 0) + group.setdefault("centered", False) + + @torch.no_grad() + def step(self, closure=None): + """Performs a single optimization step. + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + for p in group["params"]: + if p.grad is None: + continue + grad = p.grad + if grad.is_sparse: + raise RuntimeError("RMSpropTF does not support sparse gradients") + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # PyTorch initialized to zeros here + state["square_avg"] = torch.ones_like(p, memory_format=torch.preserve_format) + if group["momentum"] > 0: + state["momentum_buffer"] = torch.zeros_like(p, memory_format=torch.preserve_format) + if group["centered"]: + state["grad_avg"] = torch.zeros_like(p, memory_format=torch.preserve_format) + + square_avg = state["square_avg"] + alpha = group["alpha"] + + state["step"] += 1 + + if group["weight_decay"] != 0: + grad = grad.add(p, alpha=group["weight_decay"]) + + square_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha) + + if group["centered"]: + grad_avg = state["grad_avg"] + grad_avg.mul_(alpha).add_(grad, alpha=1 - alpha) + # PyTorch added epsilon after square root + # avg = square_avg.addcmul(grad_avg, grad_avg, value=-1).sqrt_().add_(group['eps']) + avg = square_avg.addcmul(grad_avg, grad_avg, value=-1).add_(group["eps"]).sqrt_() + else: + # PyTorch added epsilon after square root + # avg = square_avg.sqrt().add_(group['eps']) + avg = square_avg.add(group["eps"]).sqrt_() + + if group["momentum"] > 0: + buf = state["momentum_buffer"] + buf.mul_(group["momentum"]).addcdiv_(grad, avg) + p.add_(buf, alpha=-group["lr"]) + else: + p.addcdiv_(grad, avg, value=-group["lr"]) + + return loss diff --git a/stable_baselines3/common/torch_layers.py b/stable_baselines3/common/torch_layers.py index 9c74017..9429a86 100644 --- a/stable_baselines3/common/torch_layers.py +++ b/stable_baselines3/common/torch_layers.py @@ -74,7 +74,7 @@ class NatureCNN(BaseFeaturesExtractor): nn.ReLU(), nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=0), nn.ReLU(), - nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=0), + nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=0), nn.ReLU(), nn.Flatten(), ) diff --git a/stable_baselines3/common/utils.py b/stable_baselines3/common/utils.py index 9c00f0e..1c1b850 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/common/vec_env/base_vec_env.py b/stable_baselines3/common/vec_env/base_vec_env.py index 6b5a421..6338a30 100644 --- a/stable_baselines3/common/vec_env/base_vec_env.py +++ b/stable_baselines3/common/vec_env/base_vec_env.py @@ -1,5 +1,4 @@ import inspect -import pickle from abc import ABC, abstractmethod from typing import List, Optional, Sequence, Union @@ -349,7 +348,7 @@ class VecEnvWrapper(VecEnv): return shadowed_wrapper_class -class CloudpickleWrapper(object): +class CloudpickleWrapper: def __init__(self, var): """ Uses cloudpickle to serialize contents (otherwise multiprocessing tries to use pickle) @@ -362,4 +361,4 @@ class CloudpickleWrapper(object): return cloudpickle.dumps(self.var) def __setstate__(self, obs): - self.var = pickle.loads(obs) + self.var = cloudpickle.loads(obs) diff --git a/stable_baselines3/dqn/dqn.py b/stable_baselines3/dqn/dqn.py index 3d54416..ed0f655 100644 --- a/stable_baselines3/dqn/dqn.py +++ b/stable_baselines3/dqn/dqn.py @@ -8,7 +8,7 @@ from stable_baselines3.common import logger from stable_baselines3.common.buffers import ReplayBuffer 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 @@ -141,8 +141,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/dqn/policies.py b/stable_baselines3/dqn/policies.py index ea0eaa2..495b61f 100644 --- a/stable_baselines3/dqn/policies.py +++ b/stable_baselines3/dqn/policies.py @@ -133,8 +133,6 @@ class DQNPolicy(BasePolicy): else: net_arch = [] - self.features_extractor = features_extractor_class(self.observation_space, **self.features_extractor_kwargs) - self.features_dim = self.features_extractor.features_dim self.net_arch = net_arch self.activation_fn = activation_fn self.normalize_images = normalize_images @@ -142,8 +140,6 @@ class DQNPolicy(BasePolicy): self.net_args = { "observation_space": self.observation_space, "action_space": self.action_space, - "features_extractor": self.features_extractor, - "features_dim": self.features_dim, "net_arch": self.net_arch, "activation_fn": self.activation_fn, "normalize_images": normalize_images, @@ -169,7 +165,10 @@ class DQNPolicy(BasePolicy): self.optimizer = self.optimizer_class(self.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs) def make_q_net(self) -> QNetwork: - return QNetwork(**self.net_args).to(self.device) + # Make sure we always have separate networks for feature extractors etc + features_extractor = self.features_extractor_class(self.observation_space, **self.features_extractor_kwargs) + features_dim = features_extractor.features_dim + return QNetwork(features_extractor=features_extractor, features_dim=features_dim, **self.net_args).to(self.device) def forward(self, obs: th.Tensor, deterministic: bool = True) -> th.Tensor: return self._predict(obs, deterministic=deterministic) diff --git a/stable_baselines3/ppo/ppo.py b/stable_baselines3/ppo/ppo.py index eadf961..cc191d2 100644 --- a/stable_baselines3/ppo/ppo.py +++ b/stable_baselines3/ppo/ppo.py @@ -198,7 +198,7 @@ class PPO(OnPolicyAlgorithm): # Entropy loss favor exploration if entropy is None: # Approximate entropy when no analytical form - entropy_loss = -log_prob.mean() + entropy_loss = -th.mean(-log_prob) else: entropy_loss = -th.mean(entropy) diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 14d7dd7..396087b 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -9,6 +9,7 @@ from stable_baselines3.common.buffers import ReplayBuffer 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 @@ -259,8 +260,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 d7740c8..96af267 100644 --- a/stable_baselines3/td3/td3.py +++ b/stable_baselines3/td3/td3.py @@ -8,6 +8,7 @@ from stable_baselines3.common.buffers import ReplayBuffer 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 @@ -169,12 +170,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..a3df0a6 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -0.8.0a4 +0.8.0 diff --git a/tests/test_custom_policy.py b/tests/test_custom_policy.py index cd37912..c1e08df 100644 --- a/tests/test_custom_policy.py +++ b/tests/test_custom_policy.py @@ -2,6 +2,7 @@ import pytest import torch as th from stable_baselines3 import A2C, PPO, SAC, TD3 +from stable_baselines3.common.sb2_compat.rmsprop_tf_like import RMSpropTFLike @pytest.mark.parametrize( @@ -32,3 +33,8 @@ def test_custom_offpolicy(model_class, net_arch): def test_custom_optimizer(model_class, optimizer_kwargs): policy_kwargs = dict(optimizer_class=th.optim.AdamW, optimizer_kwargs=optimizer_kwargs, net_arch=[32]) _ = model_class("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(1000) + + +def test_tf_like_rmsprop_optimizer(): + policy_kwargs = dict(optimizer_class=RMSpropTFLike, net_arch=[32]) + _ = A2C("MlpPolicy", "Pendulum-v0", policy_kwargs=policy_kwargs).learn(1000) 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)