code: def calculate_gradient(self, batch_info, device, model, rollout):
          
          rollout = rollout.to_transitions()
  
          dones = rollout.batch_tensor('dones')
          rewards = rollout.batch_tensor('rewards')
          observations_next = rollout.batch_tensor('observations_next')
          actions = rollout.batch_tensor('actions')
          observations = rollout.batch_tensor('observations')
  
          # Calculate value loss - or critic loss
          with torch.no_grad():
              target_next_value = self.target_model.value(observations_next)
              target_value = rewards + (1.0 - dones) * self.discount_factor * target_next_value
  
          # Value estimation error vs the target network
          model_value = model.value(observations, actions)
          value_loss = F.mse_loss(model_value, target_value)
  
          # It may seem a bit tricky what I'm doing here, but the underlying idea is simple
          # All other implementations I found keep two separate optimizers for actor and critic
          # and update them separately
          # What I'm trying to do is to optimize them both with a single optimizer
          # but I need to make sure gradients flow correctly
          # From critic loss to critic network only and from actor loss to actor network only
  
          # Backpropagate value loss to critic only
          value_loss.backward()
  
          model_action = model.action(observations)
          model_action_value = model.value(observations, model_action)
  
          policy_loss = -model_action_value.mean()
  
          model_action_grad = torch.autograd.grad(policy_loss, model_action)[0]
  
          # Backpropagate actor loss to actor only
          model_action.backward(gradient=model_action_grad)
  
          return {
              'policy_loss': policy_loss.item(),
              'value_loss': value_loss.item(),
          }
docs: Calculate loss of the supplied rollout