Network diagram keras plot_model for tf-agents

General Discussion
, , ,
#1

Is there an example of using the keras method plot_model() for any of the tf-agents networks such as
ActorDistributionNetwork or CriticNetwork?

It doesn’t seem to work out of the box. I can kind of hack it with

actor_net._encoder._is_graph_network = True
tf.keras.utils.plot_model(
    actor_net._encoder, 'model_actor_net.png', show_shapes=True,
    show_layer_names=True, expand_nested=True
)

but this doesn’t work for the critic_net.