Running ML-Agent Trained Frozen Graph in Tensorflow or Keras

Question

I am trying to take a model trained by Unity ML-Agents and run inferences back inside Tensorflow / Keras. Any tips on how to do this?

I tried loading the frozen graph file into a TF graph and then extracting the input/output tensor. I am assuming that vector_observation is the input operation and action is the output operation.

x = graph.get_tensor_by_name('prefix/vector_observation:0')
y = graph.get_tensor_by_name('prefix/action:0')

Then I am trying to feed in the input I have on Unity (AddVectorObs)...

with tf.Session(graph=graph) as sess:
    y_out = sess.run(y, feed_dict={
            x: [[5, 5, 30, 1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36, False, False, False, False, False, False,True,False,False,False, True, False, False]]
        })

But it gave me an error about action_mask when running it.

InvalidArgumentError: You must feed a value for placeholder tensor 'prefix/action_masks' with dtype float and shape [?,6]
     [[Node: prefix/action_masks = Placeholder[dtype=DT_FLOAT, shape=[?,6], _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

My possible outputs is 2 branches like this: [0,1,2];[0,1,2]... so I guess that's why it said the action_masks has a shape of 6.

I am pretty new to Tensorflow, I apologize if this is something dumb I missed. Thanks!

Here's a picture of the neural network format:

Neural Network Input Tensor

Neural Network Architecutre