Source code for tianshou.policy.modelfree.td3

from copy import deepcopy
from typing import Any, Dict, Optional

import numpy as np
import torch

from import Batch, ReplayBuffer
from tianshou.exploration import BaseNoise, GaussianNoise
from tianshou.policy import DDPGPolicy

[docs]class TD3Policy(DDPGPolicy): """Implementation of TD3, arXiv:1802.09477. :param torch.nn.Module actor: the actor network following the rules in :class:`~tianshou.policy.BasePolicy`. (s -> logits) :param torch.optim.Optimizer actor_optim: the optimizer for actor network. :param torch.nn.Module critic1: the first critic network. (s, a -> Q(s, a)) :param torch.optim.Optimizer critic1_optim: the optimizer for the first critic network. :param torch.nn.Module critic2: the second critic network. (s, a -> Q(s, a)) :param torch.optim.Optimizer critic2_optim: the optimizer for the second critic network. :param float tau: param for soft update of the target network. Default to 0.005. :param float gamma: discount factor, in [0, 1]. Default to 0.99. :param float exploration_noise: the exploration noise, add to the action. Default to ``GaussianNoise(sigma=0.1)`` :param float policy_noise: the noise used in updating policy network. Default to 0.2. :param int update_actor_freq: the update frequency of actor network. Default to 2. :param float noise_clip: the clipping range used in updating policy network. Default to 0.5. :param bool reward_normalization: normalize the reward to Normal(0, 1). Default to False. :param bool action_scaling: whether to map actions from range [-1, 1] to range [action_spaces.low, action_spaces.high]. Default to True. :param str action_bound_method: method to bound action to range [-1, 1], can be either "clip" (for simply clipping the action) or empty string for no bounding. Default to "clip". :param Optional[gym.Space] action_space: env's action space, mandatory if you want to use option "action_scaling" or "action_bound_method". Default to None. :param lr_scheduler: a learning rate scheduler that adjusts the learning rate in optimizer in each policy.update(). Default to None (no lr_scheduler). .. seealso:: Please refer to :class:`~tianshou.policy.BasePolicy` for more detailed explanation. """ def __init__( self, actor: torch.nn.Module, actor_optim: torch.optim.Optimizer, critic1: torch.nn.Module, critic1_optim: torch.optim.Optimizer, critic2: torch.nn.Module, critic2_optim: torch.optim.Optimizer, tau: float = 0.005, gamma: float = 0.99, exploration_noise: Optional[BaseNoise] = GaussianNoise(sigma=0.1), policy_noise: float = 0.2, update_actor_freq: int = 2, noise_clip: float = 0.5, reward_normalization: bool = False, estimation_step: int = 1, **kwargs: Any, ) -> None: super().__init__( actor, actor_optim, None, None, tau, gamma, exploration_noise, reward_normalization, estimation_step, **kwargs ) self.critic1, self.critic1_old = critic1, deepcopy(critic1) self.critic1_old.eval() self.critic1_optim = critic1_optim self.critic2, self.critic2_old = critic2, deepcopy(critic2) self.critic2_old.eval() self.critic2_optim = critic2_optim self._policy_noise = policy_noise self._freq = update_actor_freq self._noise_clip = noise_clip self._cnt = 0 self._last = 0
[docs] def train(self, mode: bool = True) -> "TD3Policy": = mode self.critic1.train(mode) self.critic2.train(mode) return self
[docs] def sync_weight(self) -> None: self.soft_update(self.critic1_old, self.critic1, self.tau) self.soft_update(self.critic2_old, self.critic2, self.tau) self.soft_update(self.actor_old,, self.tau)
def _target_q(self, buffer: ReplayBuffer, indices: np.ndarray) -> torch.Tensor: batch = buffer[indices] # batch.obs: s_{t+n} act_ = self(batch, model="actor_old", input="obs_next").act noise = torch.randn(size=act_.shape, device=act_.device) * self._policy_noise if self._noise_clip > 0.0: noise = noise.clamp(-self._noise_clip, self._noise_clip) act_ += noise target_q = torch.min( self.critic1_old(batch.obs_next, act_), self.critic2_old(batch.obs_next, act_), ) return target_q
[docs] def learn(self, batch: Batch, **kwargs: Any) -> Dict[str, float]: # critic 1&2 td1, critic1_loss = self._mse_optimizer( batch, self.critic1, self.critic1_optim ) td2, critic2_loss = self._mse_optimizer( batch, self.critic2, self.critic2_optim ) batch.weight = (td1 + td2) / 2.0 # prio-buffer # actor if self._cnt % self._freq == 0: actor_loss = -self.critic1(batch.obs, self(batch, eps=0.0).act).mean() self.actor_optim.zero_grad() actor_loss.backward() self._last = actor_loss.item() self.actor_optim.step() self.sync_weight() self._cnt += 1 return { "loss/actor": self._last, "loss/critic1": critic1_loss.item(), "loss/critic2": critic2_loss.item(), }