Generator API for train_step with multiple optimizers
🚀 Feature
Add a generator API version of LightningModule.train_step() when dealing with multiple losses/optimizers.
Motivation
When dealing with multiple optimizers, the current API pass an additional optimizer_idx argument to train_step(). This is cumbersome with optional steps and shared variables. For example, this code section
        self.q1_network_optimizer.zero_grad()
        q1_value = self.q1_network(state, action)
        q1_loss = self.q_network_loss(q1_value, target_q_value)
        q1_loss.backward()
        self.q1_network_optimizer.step()

        if self.q2_network:
            self.q2_network_optimizer.zero_grad()
            q2_value = self.q2_network(state, action)
            q2_loss = self.q_network_loss(q2_value, target_q_value)
            q2_loss.backward()
            self.q2_network_optimizer.step()

would have to stash target_q_value into a member variable. It would also need to have a member variable to track the optimizer index of q2_network_optimizer, e.g.,
def __init__(self, ...):
  self.q2_optimizer_index = 1 if self.q2_network else None
  self.actor_network_optimizer = 2 if self.q2_network else 1
  self.target_q_value = None

def train_step(self, batch, batch_idx, optimizer_idx):
   if self.target_q_value is None:
    self.target_q_value = self.compute_target_q_value(...)

  if optimizer_idx == 1:
    q1_loss = self.q1_loss(self.target_q_value, ...)
    return q1_loss

  if optimizer_idx == self.q2_optimizer_index:
     q2_loss = self.q2_loss(self.target_q_value, ...)
     return q2_loss

  # reset self.target_q_value to None if this is the last optimizer

  if optimizer_idx == self.actor_network_optimizer:
     ...

Pitch
With a generator API, it would be very natural
        q1_value = self.q1_network(state, action)
        q1_loss = self.q_network_loss(q1_value, target_q_value)
        yield q1_loss

        if self.q2_network:
            q2_value = self.q2_network(state, action)
            q2_loss = self.q_network_loss(q2_value, target_q_value)
            yield q2_loss

This is much more readable and would be a lot less error-prone
Additional context