Issue with Optimizer Update in A2C Network with Optax Body:

Hello everyone,

I've encountered a problem while implementing an A2C (Advantage Actor-Critic) network involving Flax and Optax. My network includes _policy_network_ and _value_network_, each containing _policy_head_ and _torso_. When attempting to use _optimizer.update(grad)_, I received the following error:

_ValueError: Mismatch custom node data: ('policy_head', 'torso') != ('policy_network', 'value_network');_

The error message indicates that the expected keys are (_'policy_network', 'value_network')_, but the actual provided keys are ('_policy_head', 'torso'_). The structure of my model parameters is as follows:

_State({
'policy_network': {
'policy_head': {...},
'torso': {...},
},
'value_network': {
'policy_head': {...},
'torso': {...},
})_
I have tried to combine the model parameters and pass them to the optimizer, like this:

_params = {'w1': model1_params, 'w2': model2_params}_

However, this approach did not resolve the issue. I'm wondering if there is another way to correctly initialize and update the parameters of the A2C network's parameters using Optax in Flax.

If you have any suggestions or need more information, please let me know. Thank you very much for your help!