Network diagram keras plot_model for tf-agents

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
    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.