convert stable-baselines tensorflow model to tensorflowjs

Question

I am trying to convert the underlying tensorflow model from stable-baselines to tensorflowjs to be able to use the model on the browser. But I could not make the conversion work

I followed this github issue to create the necessary tensorflow files using the code:

def generate_checkpoint_from_model(model, checkpoint_name):  
        tf.saved_model.simple_save(model.sess, checkpoint_name, inputs={"obs": model.act_model.obs_ph}, outputs={"action": model.action_ph})

Then I try to transform the model using tensorflowjs_converter

tensorflowjs_converter --input_format=tf_saved_model test/ web_test

However, it gives me the following error:

Unable to lift tensor <tf.Tensor 'loss/action_ph:0' shape=(?,) dtype=int32> because it depends transitively on placeholder <tf.Operation 'loss/action_ph' type=Placeholder> via at least one path, e.g.: loss/action_ph (Placeholder)

I created the following colab notebook with the error so you can try it.

Does anyone knows how to make this conversion work?

Thank you for the help

Answer 1

I posted the question as an issue in stable-baselines and they answered. I will copy here as a reference to others:

You are trying to save the action placeholder used in PPO training (part of PPO agent), but for inference you only need the trained policy and its placeholders (model.act_model). The code on colab runs without errors by changing call to simple_save to this:

tf.saved_model.simple_save(model.sess, checkpoint_name, inputs={"obs": model.act_model.obs_ph}, outputs={"action": model.act_model._policy_proba})

The value of _policy_proba depends on the environment/algorithm.