I have a NN defined in pytorch and I have created two instances of that net as self.actor_critic_r1 and self.actor_critic_r2. I calculate the losses of each net i.e. loss1 and loss2 and I sum it up and calculate the grads in the following way,
loss_r1 = value_loss_r1 + action_loss_r1 - dist_entropy_r1 * args.entropy_coef
loss_r2 = value_loss_r2 + action_loss_r2 - dist_entropy_r2 * args.entropy_coef
self.optimizer_r1.zero_grad()
self.optimizer_r2.zero_grad()
loss = loss_r1 + loss_r2
loss.backward()
self.optimizer_r1.step()
self.optimizer_r2.step()
clip_grad_norm_(self.actor_critic_r1.parameters(), args.max_grad_norm)
clip_grad_norm_(self.actor_critic_r2.parameters(), args.max_grad_norm)
Alternatively, should I update the loss individually like this,
self.optimizer_r1.zero_grad()
(value_loss_r1 + action_loss_r1 - dist_entropy_r1 * args.entropy_coef).backward()
self.optimizer_r1.step()
clip_grad_norm_(self.actor_critic_r1.parameters(), args.max_grad_norm)
self.optimizer_r2.zero_grad()
(value_loss_r2 + action_loss_r2 - dist_entropy_r2 * args.entropy_coef).backward()
self.optimizer_r2.step()
clip_grad_norm_(self.actor_critic_r2.parameters(), args.max_grad_norm)
I am not sure if this the right approach to update a network with multiple loss please provide your suggestion.