r"""

config = dict(

# (str) RL policy register name (refer to function "POLICY_REGISTRY").

type='sac',

# (bool) Whether to use cuda for network.

cuda=True,

# (bool type) on_policy: Determine whether on-policy or off-policy.

# on-policy setting influences the behaviour of buffer.

# Default False in SAC.

on_policy=False,

# (bool type) priority: Determine whether to use priority in buffer sample.

# Default False in SAC.

priority=False,

# (bool) Whether use Importance Sampling Weight to correct biased update. If True, priority must be True.

priority_IS_weight=False,

# (int) Number of training samples(randomly collected) in replay buffer when training starts.

# Default 10000 in SAC.

random_collect_size=10000,

model=dict(

# (bool type) twin_critic: Determine whether to use double-soft-q-net for target q computation.

# Please refer to TD3 about Clipped Double-Q Learning trick, which learns two Q-functions instead of one .

# Default to True.

twin_critic=True,

# (bool type) value_network: Determine whether to use value network as the

# original SAC paper (arXiv 1801.01290).

# using value_network needs to set learning_rate_value, learning_rate_q,

# and learning_rate_policy in `cfg.policy.learn`.

# Default to False.

# value_network=False,

actor_head_type='reparameterization',

),

learn=dict(

# (bool) Whether to use multi gpu

multi_gpu=False,

# How many updates(iterations) to train after collector's one collection.

# Bigger "update_per_collect" means bigger off-policy.

# collect data -> update policy-> collect data -> ...

update_per_collect=1,

# (int) Minibatch size for gradient descent.

batch_size=256,

# (float type) learning_rate_q: Learning rate for soft q network.

# Default to 3e-4.

# Please set to 1e-3, when model.value_network is True.

learning_rate_q=3e-4,

# (float type) learning_rate_policy: Learning rate for policy network.

# Default to 3e-4.

# Please set to 1e-3, when model.value_network is True.

learning_rate_policy=3e-4,

# (float type) learning_rate_value: Learning rate for value network.

# `learning_rate_value` should be initialized, when model.value_network is True.

# Please set to 3e-4, when model.value_network is True.

learning_rate_value=3e-4,

# (float type) learning_rate_alpha: Learning rate for auto temperature parameter `\alpha`.

# Default to 3e-4.

learning_rate_alpha=3e-4,

# (float type) target_theta: Used for soft update of the target network,

# aka. Interpolation factor in polyak averaging for target networks.

# Default to 0.005.

target_theta=0.005,

# (float) discount factor for the discounted sum of rewards, aka. gamma.

discount_factor=0.99,

# (float type) alpha: Entropy regularization coefficient.

# Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details.

# If auto_alpha is set to `True`, alpha is initialization for auto `\alpha`.

# Default to 0.2.

alpha=0.2,

# (bool type) auto_alpha: Determine whether to use auto temperature parameter `\alpha` .

# Temperature parameter determines the relative importance of the entropy term against the reward.

# Please check out the original SAC paper (arXiv 1801.01290): Eq 1 for more details.

# Default to False.

# Note that: Using auto alpha needs to set learning_rate_alpha in `cfg.policy.learn`.

auto_alpha=True,

# (bool type) log_space: Determine whether to use auto `\alpha` in log space.

log_space=True,

# (bool) Whether ignore done(usually for max step termination env. e.g. pendulum)

# Note: Gym wraps the MuJoCo envs by default with TimeLimit environment wrappers.

# These limit HalfCheetah, and several other MuJoCo envs, to max length of 1000.

# However, interaction with HalfCheetah always gets done with False,

# Since we inplace done==True with done==False to keep

# TD-error accurate computation(``gamma * (1 - done) * next_v + reward``),

# when the episode step is greater than max episode step.

ignore_done=False,

# (float) Weight uniform initialization range in the last output layer

init_w=3e-3,

),

collect=dict(

# You can use either "n_sample" or "n_episode" in actor.collect.

# Get "n_sample" samples per collect.

# Default n_sample to 1.

n_sample=1,

# (int) Cut trajectories into pieces with length "unroll_len".

unroll_len=1,

),

eval=dict(

evaluator=dict(

# (int) Evaluate every "eval_freq" training iterations.

eval_freq=5000,

),

),

other=dict(

replay_buffer=dict(

# (int type) replay_buffer_size: Max size of replay buffer.

replay_buffer_size=1000000,

# (int type) max_use: Max use times of one data in the buffer.

# Data will be removed once used for too many times.

# Default to infinite.

# max_use=256,

),

),

)

def _init_learn(self) -> None:

r"""

# Init

self._priority = self._cfg.priority

self._priority_IS_weight = self._cfg.priority_IS_weight

self._value_network = False # TODO self._cfg.model.value_network

self._twin_critic = self._cfg.model.twin_critic

# Weight Init for the last output layer

init_w = self._cfg.learn.init_w

self._model.actor[2].mu.weight.data.uniform_(-init_w, init_w)

self._model.actor[2].mu.bias.data.uniform_(-init_w, init_w)

self._model.actor[2].log_sigma_layer.weight.data.uniform_(-init_w, init_w)

self._model.actor[2].log_sigma_layer.bias.data.uniform_(-init_w, init_w)

# Optimizers

if self._value_network:

self._optimizer_value = Adam(

self._model.value_critic.parameters(),

lr=self._cfg.learn.learning_rate_value,

)

self._optimizer_q = Adam(

self._model.critic.parameters(),

lr=self._cfg.learn.learning_rate_q,

)

self._optimizer_policy = Adam(

self._model.actor.parameters(),

lr=self._cfg.learn.learning_rate_policy,

)

# Algorithm config

self._gamma = self._cfg.learn.discount_factor

# Init auto alpha

if self._cfg.learn.auto_alpha:

self._target_entropy = self._cfg.learn.get('target_entropy', -np.prod(self._cfg.model.action_shape))

if self._cfg.learn.log_space:

self._log_alpha = torch.log(torch.FloatTensor([self._cfg.learn.alpha]))

self._log_alpha = self._log_alpha.to(self._device).requires_grad_()

self._alpha_optim = torch.optim.Adam([self._log_alpha], lr=self._cfg.learn.learning_rate_alpha)

assert self._log_alpha.shape == torch.Size([1]) and self._log_alpha.requires_grad

self._alpha = self._log_alpha.detach().exp()

self._auto_alpha = True

self._log_space = True

else:

self._alpha = torch.FloatTensor([self._cfg.learn.alpha]).to(self._device).requires_grad_()

self._alpha_optim = torch.optim.Adam([self._alpha], lr=self._cfg.learn.learning_rate_alpha)

self._auto_alpha = True

self._log_space = False

else:

self._alpha = torch.tensor(

[self._cfg.learn.alpha], requires_grad=False, device=self._device, dtype=torch.float32

)

self._auto_alpha = False

# Main and target models

self._target_model = copy.deepcopy(self._model)

self._target_model = model_wrap(

self._target_model,

wrapper_name='target',

update_type='momentum',

update_kwargs={'theta': self._cfg.learn.target_theta}

)

self._learn_model = model_wrap(self._model, wrapper_name='base')

self._learn_model.reset()

self._target_model.reset()

self._forward_learn_cnt = 0

if self._cfg.learn.get("sphere_optim", None):

self.q_fc_layer=[m for m in self._model.critic.modules() if isinstance(m, torch.nn.Linear)]

self.init_q_fc_norm=self._compute_norm([l.weight.clone().detach() for l in self.q_fc_layer])

def _forward_learn(self, data: dict) -> Dict[str, Any]:

"""

loss_dict = {}

data = default_preprocess_learn(

data,

use_priority=self._priority,

use_priority_IS_weight=self._cfg.priority_IS_weight,

ignore_done=self._cfg.learn.ignore_done,

use_nstep=False

)

if self._cuda:

data = to_device(data, self._device)

self._learn_model.train()

self._target_model.train()

obs = data['obs']

next_obs = data['next_obs']

reward = data['reward']

done = data['done']

# 1. predict q value

q_value = self._learn_model.forward(data, mode='compute_critic')['q_value']

# 2. predict target value depend self._value_network.

if self._value_network:

v_value = self._learn_model.forward(obs, mode='compute_value_critic')['v_value']

with torch.no_grad():

next_v_value = self._target_model.forward(next_obs, mode='compute_value_critic')['v_value']

target_q_value = next_v_value

else:

# target q value.

with torch.no_grad():

(mu, sigma) = self._learn_model.forward(next_obs, mode='compute_actor')['logit']

dist = Independent(Normal(mu, sigma), 1)

pred = dist.rsample()

next_action = torch.tanh(pred)

y = 1 - next_action.pow(2) + 1e-6

# keep dimension for loss computation (usually for action space is 1 env. e.g. pendulum)

next_log_prob = dist.log_prob(pred).unsqueeze(-1)

next_log_prob = next_log_prob - torch.log(y).sum(-1, keepdim=True)

next_data = {'obs': next_obs, 'action': next_action}

target_q_value = self._target_model.forward(next_data, mode='compute_critic')['q_value']

# the value of a policy according to the maximum entropy objective

if self._twin_critic:

# find min one as target q value

target_q_value = torch.min(target_q_value[0],

target_q_value[1]) - self._alpha * next_log_prob.squeeze(-1)

# target_q_value = torch.cat([i.unsqueeze(-1) for i in target_q_value], dim=-1) - self._alpha * next_log_prob

else:

target_q_value = target_q_value - self._alpha * next_log_prob.squeeze(-1)

# 3. compute q loss

if self._twin_critic:

q_data0 = v_1step_td_data(q_value[0], target_q_value, reward, done, data['weight'])

loss_dict['critic_loss'], td_error_per_sample0 = v_1step_td_error(q_data0, self._gamma)

q_data1 = v_1step_td_data(q_value[1], target_q_value, reward, done, data['weight'])

loss_dict['twin_critic_loss'], td_error_per_sample1 = v_1step_td_error(q_data1, self._gamma)

td_error_per_sample = (td_error_per_sample0 + td_error_per_sample1) / 2

# q_value=torch.cat([i.unsqueeze(-1) for i in q_value], dim=-1)

# td_error_per_sample = torch.nn.MSELoss(reduction='none')(q_value, reward.view(-1,1) + self._gamma * target_q_value)

# td_error_per_sample = td_error_per_sample.mean()

# loss_dict['critic_loss'] = td_error_per_sample

# loss_dict['twin_critic_loss']=0.0

else:

q_data = v_1step_td_data(q_value, target_q_value, reward, done, data['weight'])

loss_dict['critic_loss'], td_error_per_sample = v_1step_td_error(q_data, self._gamma)

# 4. update q network

self._optimizer_q.zero_grad()

loss_dict['critic_loss'].backward()

if self._twin_critic:

loss_dict['twin_critic_loss'].backward()

self._optimizer_q.step()

# sphere optimization

if self._cfg.learn.get("sphere_optim", None):

curr_weight_norm = self._compute_norm([l.weight.clone().detach() for l in self.q_fc_layer])

for l in self.q_fc_layer:

l.weight.data.mul_(self.init_q_fc_norm / curr_weight_norm)

# 5. evaluate to get action distribution

(mu, sigma) = self._learn_model.forward(data['obs'], mode='compute_actor')['logit']

dist = Independent(Normal(mu, sigma), 1)

pred = dist.rsample()

action = torch.tanh(pred)

y = 1 - action.pow(2) + 1e-6

# keep dimension for loss computation (usually for action space is 1 env. e.g. pendulum)

log_prob = dist.log_prob(pred).unsqueeze(-1)

log_prob = log_prob - torch.log(y).sum(-1, keepdim=True)

eval_data = {'obs': obs, 'action': action}

new_q_value = self._learn_model.forward(eval_data, mode='compute_critic')['q_value']

if self._twin_critic:

new_q_value = torch.min(new_q_value[0], new_q_value[1])

# 6. (optional) compute value loss and update value network

if self._value_network:

# new_q_value: (bs, ), log_prob: (bs, act_shape) -> target_v_value: (bs, )

target_v_value = (new_q_value.unsqueeze(-1) - self._alpha * log_prob).mean(dim=-1)

loss_dict['value_loss'] = F.mse_loss(v_value, target_v_value.detach())

# update value network

self._optimizer_value.zero_grad()

loss_dict['value_loss'].backward()

self._optimizer_value.step()

# 7. compute policy loss

policy_loss = (self._alpha * log_prob - new_q_value.unsqueeze(-1)).mean()

loss_dict['policy_loss'] = policy_loss

# 8. update policy network

self._optimizer_policy.zero_grad()

loss_dict['policy_loss'].backward()

self._optimizer_policy.step()

# 9. compute alpha loss

if self._auto_alpha:

if self._log_space:

log_prob = log_prob + self._target_entropy

loss_dict['alpha_loss'] = -(self._log_alpha * log_prob.detach()).mean()

self._alpha_optim.zero_grad()

loss_dict['alpha_loss'].backward()

self._alpha_optim.step()

self._alpha = self._log_alpha.detach().exp()

else:

log_prob = log_prob + self._target_entropy

loss_dict['alpha_loss'] = -(self._alpha * log_prob.detach()).mean()

self._alpha_optim.zero_grad()

loss_dict['alpha_loss'].backward()

self._alpha_optim.step()

self._alpha = max(0, self._alpha)

loss_dict['total_loss'] = sum(loss_dict.values())

# =============

# after update

# =============

self._forward_learn_cnt += 1

# target update

self._target_model.update(self._learn_model.state_dict())

return {

'cur_lr_q': self._optimizer_q.defaults['lr'],

'cur_lr_p': self._optimizer_policy.defaults['lr'],

'q_fc_norm': self._get_fc_weight_norm(self._model.critic),

'policy_fc_norm': self._get_fc_weight_norm(self._model.actor),

'priority': td_error_per_sample.abs().tolist(),

'td_error': td_error_per_sample.detach().mean().item(),

'alpha': self._alpha.item(),

'target_q_value': target_q_value.detach().mean().item(),

**loss_dict

}

def _get_fc_weight_norm(self, net):

with torch.no_grad():

return torch.sqrt(sum([torch.sum(m.weight.clone()**2) for m in net.modules() if isinstance(m, torch.nn.Linear)]))

def _compute_norm(self, tensor_list):

sq_sum = 0

for t in tensor_list:

sq_sum += torch.sum(t ** 2)

return torch.sqrt(sq_sum)

def _state_dict_learn(self) -> Dict[str, Any]:

ret = {

'model': self._learn_model.state_dict(),

'target_model': self._target_model.state_dict(),

'optimizer_q': self._optimizer_q.state_dict(),

'optimizer_policy': self._optimizer_policy.state_dict(),

}

if self._value_network:

ret.update({'optimizer_value': self._optimizer_value.state_dict()})

if self._auto_alpha:

ret.update({'optimizer_alpha': self._alpha_optim.state_dict()})

return ret

def _load_state_dict_learn(self, state_dict: Dict[str, Any]) -> None:

self._learn_model.load_state_dict(state_dict['model'])

self._target_model.load_state_dict(state_dict['target_model'])

self._optimizer_q.load_state_dict(state_dict['optimizer_q'])

if self._value_network:

self._optimizer_value.load_state_dict(state_dict['optimizer_value'])

self._optimizer_policy.load_state_dict(state_dict['optimizer_policy'])

if self._auto_alpha:

self._alpha_optim.load_state_dict(state_dict['optimizer_alpha'])

def _init_collect(self) -> None:

r"""

self._unroll_len = self._cfg.collect.unroll_len

self._collect_model = model_wrap(self._model, wrapper_name='base')

self._collect_model.reset()

def _forward_collect(self, data: dict) -> dict:

r"""

data_id = list(data.keys())

data = default_collate(list(data.values()))

if self._cuda:

data = to_device(data, self._device)

self._collect_model.eval()

with torch.no_grad():

(mu, sigma) = self._collect_model.forward(data, mode='compute_actor')['logit']

dist = Independent(Normal(mu, sigma), 1)

action = torch.tanh(dist.rsample())

output = {'logit': (mu, sigma), 'action': action}

if self._cuda:

output = to_device(output, 'cpu')

output = default_decollate(output)

return {i: d for i, d in zip(data_id, output)}

def _process_transition(self, obs: Any, policy_output: dict, timestep: namedtuple) -> dict:

r"""

transition = {

'obs': obs,

'next_obs': timestep.obs,

'action': policy_output['action'],

'reward': timestep.reward,

'done': timestep.done,

}

return transition

def _get_train_sample(self, data: list) -> Union[None, List[Any]]:

return get_train_sample(data, self._unroll_len)

def _init_eval(self) -> None:

r"""

self._eval_model = model_wrap(self._model, wrapper_name='base')

self._eval_model.reset()

def _forward_eval(self, data: dict) -> dict:

r"""

data_id = list(data.keys())

data = default_collate(list(data.values()))

if self._cuda:

data = to_device(data, self._device)

self._eval_model.eval()

with torch.no_grad():

(mu, sigma) = self._eval_model.forward(data, mode='compute_actor')['logit']

action = torch.tanh(mu) # deterministic_eval

output = {'action': action}

if self._cuda:

output = to_device(output, 'cpu')

output = default_decollate(output)

return {i: d for i, d in zip(data_id, output)}

def default_model(self) -> Tuple[str, List[str]]:

return 'qac', ['ding.model.template.qac']

def _monitor_vars_learn(self) -> List[str]:

r"""

twin_critic = ['twin_critic_loss'] if self._twin_critic else []

alpha_loss = ['alpha_loss'] if self._auto_alpha else []

value_loss = ['value_loss'] if self._value_network else []

return [

'alpha_loss',

'policy_loss',

'critic_loss',

'cur_lr_q',

'cur_lr_p',

'target_q_value',

'alpha',

'td_error',

'q_fc_norm',

'policy_fc_norm'