Add another apply_fn in TrainState

[vcharraut]

Hello!

I have first implemented an A2C algorithm in Flax here but for the ease of the implementation I use a shared network for the Actor Critic model.

What I would like to do is define two model as follow:

class ActorNet(nn.Module):
    num_actions: int

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(features=64)(x)
        x = nn.tanh(x)
        x = nn.Dense(features=64)(x)
        x = nn.tanh(x)

        logits = nn.Dense(features=self.num_actions, name="actor")(x)
        log_probs = nn.log_softmax(logits)

        return log_probs


class CriticNet(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(features=64)(x)
        x = nn.tanh(x)
        x = nn.Dense(features=64)(x)
        x = nn.tanh(x)

        return nn.Dense(features=1, name="critic")(x).squeeze()

but only use one TrainState since the two models share the same data and optimizer.

I guess the right way of doing so is to create a custom TrainState class but I'm not sure how to define two distinct apply_fn.

class TrainState(TrainState):
    critic_params: flax.core.FrozenDict
    apply_critic: nn.Module.apply


train_state = TrainState.create(
    apply_fn=actor.apply,
    params=actor_params,
    critic_params=critic_params,
    apply_critic=critic.apply,
    tx=optimizer
)

I know this code is not working but I would be happy if someone could guide me towards the right solution

Thanks

[cgarciae]

Hey @valentin-cnt, some notes:

• If you are planning on using TrainState.apply_gradients they would only be applied to params=actor_params, not critic_params.
• You need to mark apply_critic as a "static" field since its not a valid JAX type:

from flax.struct import field

class TrainState(TrainState):
    critic_params: flax.core.FrozenDict
    apply_critic: Callable = field(pytree_node=False)

[vcharraut]

If you are planning on using TrainState.apply_gradients they would only be applied to params=actor_params, not critic_params.

Ok, actually I didn't think about this step. I guess it would become a bad pratice to redefine TrainState.apply_gradients, I rather not touch that
Defining every operations in the one combined ActorCriticNet seems to be the better solution

Thank you for your answer!