Best way to export a keras model from RLlib

Hi all! I have executed many training iterations with a PPO agent and now I want to export ONLY the underlying keras model that contains the trained neural network, so that I can eveluate its performance outside RLlib (only as a keras model with model.predict()) Is there any way to export only this part of the Tensorflow graph (maybe as a h5 file). I want to get something that later can be restored to a keras model with tf.keras.model.load_model('model.h5').

Thanks in advance!

Have you tried:

Trainer.get_policy().export_model([export dir])

?
This will create a tf1.x snapshot of the model’s tf static graph and you should be able to re-load it later into a new session (just the model).

The problem by doing so was that the whole TF Graph was saved, and it contained many nodes that were useless for my purpose… What I wanted to export was only the part of the graph containing the keras model, in order to be able to load it later with tf.keras.models.load_model('model.h5').

I tried to achieve this by doing:

agent = ppo.PPOTrainer(config, env='Pong-v0')
agent.restore(checkpoint_dir)

with agent.get_policy().get_session().graph.as_default():
    export_model = agent.get_policy().model.base_model.save('model.h5')

but the weights saved seemed to be different form those that the model have (and that were accesible via agent.get_policy().get_weights() function).