I don’t think it necessarily is a problem, but if you disregard the update between step 1 and step 2 and the potentially different optimizer hyperparameters, you are effectively optimizing the shared weights to a sum of the two objective functions

True, true! I was just wondering if this might pose a problem since, as we both mentioned, it’s optimizing over two objectives (hence my scaling the gradient update for the shared layers).

P.S.: Is it intentional that you take the second output of actor_critic for both?

Yeah, the first output is the output of the actor, and the second is the output of the critic. Our actor_critic has 2 “cases”:


self.actor_critic(state, action), the critic head uses the passed-in action to output the Q-value. (line 13, the LHS of the optimization eqn)

Q_{\phi}(s, a) where (s, a) are sampled from the replay buffer.


self.actor_critic(state) where we do not pass in an action. Here, we use our policy head to output a policy and we then output the Q-value from that (Note how in line 14, of the algorithm, we are taking the derivative of the Q-value but WRT the actor’s parameters)

Q_{\phi}(s, mu_{\theta}(a)). I needed to be careful with how I generated the gradient but I think I did it correctly.
The below is the code implementing “case 2”
Optimizer
params = list(self.shared_module.parameters()) + list(self.mu.parameters())
self.actor_optimizer = optim.Adam(params, lr=self.actor_lr)

Forward pass in “case 2”
compressed: T.Tensor = self.shared_layer(state)
########################################
# Actor runs first
########################################
mu_val = self.mu(F.relu(compressed))
actions: T.Tensor = T.tanh(mu_val)
leq, geq = actions <= 1, actions >= -1
if not torch.all(leq * geq):
    raise Exception("Our actions were out of the expected range")

########################################
# Critic runs now
########################################
action_value = F.relu(self.action_value(actions))
state_action_value = self.q(
    F.relu(T.add(compressed, action_value))
)
return actions, state_action_value

And the update step is as in my original post.

Does anything look egregiously wrong?