DDPG actor-critic with shared layer?

mike_and_ike: <p>Thanks I’m implementing DDPG where the actor and critic have a shared module. I’m running into an issue and I was wondering if I could get some feedback. I have the following:</p>
<pre><code class="lang-auto">INPUT_DIM = 100
BOTTLENECK_DIMS = 10
class SharedModule(nn.Module):
  def __init__(self):
    self.shared = nn.Linear(INPUT_DIM, BOTTLENECK_DIMS)
  def forward(self, x):
    return self.shared(x)

class ActorCritic(nn.Module):
  
  def __init__(self, n_actions, shared: SharedModule):
    self.shared = shared
    self.n_actions = n_actions
    
    # Critic definition 
    self.action_value = nn.Linear(self.n_actions, BOTTLENECK_DIMS)
    self.q = nn.Linear(BOTTLENECK_DIMS, 1)

    # Actor Definition
    self.mu = nn.Linear(BOTTLENECK_DIMS, self.n_actions)

    self.optimizer = optim.Adam(self.parameters(), lr=self.lr)

  def forward(self, state, optional_action=None):
    if optional_action:
        return self._wo_action_fwd(state)
    return self._w_action_fwd(state, optional_action)
    
  def _wo_action_fwd(self, state):
    shared_output = self.shared(state)

    # Computing the actions
    mu_val = self.mu(F.relu(shared_output))
    actions = T.tanh(mu_val)

    # Computing the Q-vals
    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

  def _w_action_forward(self, state, action):
    shared_output = self.shared(state)
    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
</code></pre>
<p>My training process is then</p>
<pre><code class="lang-auto">shared_module = SharedModule()
actor_critic = ActorCritic(n_actions=3, shared_module)

shared_module = SharedModule()
T_actor_critic = ActorCritic(n_actions=3, shared_module)

s_batch, a_batch, r_batch, s_next_batch, d_batch = memory.sample(batch_size)

########################################
# Generate labels
########################################
# Get our critic target
_, y_critic = T_actor_critic(s_next_batch)
target = T.unsqueeze(
    r_batch + (gamma * d_batch * T.squeeze(y_critic)),
    dim=-1
)

########################################
# Critic Train
########################################
actor_critic.optimizer.zero_grad()
_, y_hat_critic = actor_critic(s_batch, a_batch)
critic_loss = F.mse_loss(target, y_hat_critic)
critic_loss.backward()
actor_critic.optimizer.step()

########################################
# Actor train
########################################
actor_critic.optimizer.zero_grad()
_, y_hat_policy = actor_critic(s_batch)
policy_loss = T.mean(-y_hat_policy)
policy_loss.backward()
actor_critic.optimizer.step()
</code></pre>
<h1>
<a name="issues-doubts-1" class="anchor" href="#issues-doubts-1"></a>Issues / doubts</h1>
<ol>
<li>Looking at <a href="https://spinningup.openai.com/en/latest/_images/math/5811066e89799e65be299ec407846103fcf1f746.svg" rel="noopener nofollow ugc">OpenAI DDPG formula discussion</a>, I’ve done step 12 and step 13 correctly (as far as I can tell). However, I don’t know how to do step 14.</li>
</ol>
<p>The issue is that although I can calculate the entire Q-value, I don’t know how to take the derivative only with regards to theta. How should I go about doing this? I tried using</p>
<pre><code class="lang-auto">  def _wo_action_fwd(self, state):
    shared_output = self.shared(state)

    # Computing the actions
    mu_val = self.mu(F.relu(shared_output))
    actions = T.tanh(mu_val)

    # Computing the Q-vals
    with T.no_grad():
      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
</code></pre>
<p>but unsurprisingly I got an error of <code>has no gradients</code> which is unsurprising since I dropped off the gradients.</p>
<ol start="2">
<li>This is more of a DDPG question as opposed to a pytorch one, but is my translation of the algorithm correct? I do a step for the critic and then one for the actor?  I’ve seen</li>
</ol>
<p><code>loss = torch.stack(policy_losses).sum() + torch.stack(value_losses).sum()</code></p>
<ol start="3">
<li>Is there a way to train it so that the shared module is stable? I imagine that being trained on two separate losses (I’m optimizing over 2 steps) might make convergence of that shared module wonky.</li>
</ol>

mike_and_ike: <p>In response to <span class="hashtag">#2</span>, I found that line of code from <a href="https://github.com/pytorch/examples/blob/master/reinforcement_learning/actor_critic.py#L126" class="inline-onebox" rel="noopener nofollow ugc">examples/actor_critic.py at master · pytorch/examples · GitHub</a> which I found from <a href="https://discuss.pytorch.org/t/actor-critic-implementation-problem/12782" class="inline-onebox">Actor Critic implementation problem</a></p>
<p>Since I’m new I can’t post more than 2 links in my post womp womp</p>