Using custom keras model with layer sharing together dqn_agent.DqnAgent()

Question

I am trying to use a custom neural network with the DqnAgent() from tf. In my model I need to use layer sharing. Thus, I use the functional API to build the model. The model has a dict as input and one layer with n neurons as output. The last layer is a Concat- and not a Dens-Layer though. The type of the model that i get from the functional API keras.Model(inputs=[...], outpunts[...] is "keras.engine.functional.Functional".

Now i want to use my Model with the tf Agent like this:

agent = dqn_agent.DqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    q_network=model,
    optimizer=optimizer,
    td_errors_loss_fn=common.element_wise_squared_loss,
    train_step_counter=train_step_counter)

I get the following Error thoug:

AttributeError: 'Functional' object has no attribute 'create_variables'
  In call to configurable 'DqnAgent' (<class 'tf_agents.agents.dqn.dqn_agent.DqnAgent'>)

The q_network expects a network from type "network.Network". I am not sure how to convert or wrap my environment in such a way, that the DqnAgent() will accept it. How could I manage to do this? Any support is much appreciated. If you need more information about anything let me know.

Additional information about my Network:

• Input dict with multiple inputs.
• multiple shared dens layers. output of the last one is shape (1,).
• concatenate all those outputs of shape (1,)
• one multipy layer to eliminate infeasible actions by multiplying the outputs with 0 or 1 respectively.