How to use PPOLoss with shared actor and critic parameters?

kf064: <p>I am training an actor-critic model in a reinforcement learning environment which processes observations in a shared hidden module before passing the data to the policy and value heads <a href="https://i.stack.imgur.com/F4kJv.png" rel="noopener nofollow ugc">as shown in this diagram.</a></p>
<p>The PyTorch documentation for ClipPPOLoss indicates that you can train these kinds of networks in the following way:</p>
<blockquote>
<p>If the actor and the value function share parameters, one can avoid calling the common module multiple times by passing only the head of the value network to the PPO loss module:</p>
</blockquote>
<pre><code class="lang-auto">common = SomeModule(in_keys=["observation"], out_keys=["hidden"])
actor_head = SomeActor(in_keys=["hidden"])
value_head = SomeValue(in_keys=["hidden"])
# first option, with 2 calls on the common module
model = ActorCriticOperator(common, actor_head, value_head)
loss_module = PPOLoss(model.get_policy_operator(), model.get_value_operator())
# second option, with a single call to the common module
loss_module = PPOLoss(ProbabilisticTensorDictSequential(model, actor_head), value_head)
</code></pre>
<p>I am trying to implement the second option, using <a href="https://pytorch.org/rl/_modules/torchrl/objectives/value/advantages.html#GAE" rel="noopener nofollow ugc">TorchRL’s Generalized Advantage Estimation</a> as the <code>loss_module.value_estimator</code>. Using the <code>model</code> as defined above works fine for filling up a <code>SyncDataCollector</code> with frames, but I’m running into issues with calculating the advantage on each training iteration, after the batch has been sampled.</p>
<pre><code class="lang-auto">with torch.no_grad():
    GAE(
            tensordict_data,
            params=loss_module.critic_network_params,
            target_params=loss_module.target_critic_network_params,
        ) 
</code></pre>
<p>Even though the collected <code>tensordict_data</code> contains the <code>hidden</code> key, <code>value_head</code> is being fed this empty input:</p>
<pre><code class="lang-auto">(NonTensorData(
    data=None,
    batch_size=torch.Size([]),
    device=None,
    is_shared=False),)
</code></pre>
<p>I’ve traced the error to GAE calling the <code>_call_value_nets()</code> function in the linked documentation above; this function is searching for the <code>"hidden"</code> key in the current frame (which is present) as well as in the nested <code>("next", "hidden")</code> value. This produces the empty input because the <code>"hidden"</code> key has not been defined in <code>"next"</code>. I tried to get around this by setting the “target parameters” to <code>None</code>, and <code>shifted=True</code>:</p>
<pre><code class="lang-auto">loss_module.make_value_estimator(
    ValueEstimators.GAE, gamma=gamma, lmbda=lmbda, shifted=True
)

GAE(
   tensordict_data,
   params=loss_module.critic_network_params,
   target_params=None,
   ) 
</code></pre>
<p>But this throws the error</p>
<pre><code class="lang-auto"> KeyError: "got keys {'hidden', 'state_value'} and set() which are incompatible 
</code></pre>
<p>Is there a way to define <code>hidden</code> in <code>next</code>, or is there something that I’m missing to implement PPO with shared actor/critic parameters?</p>