Building a reinforcement learning agent with JAX, and deploying it on Android with TensorFlow Lite

# Convert to tflite model

 model = PolicyGradient()

 jax_predict_fn = lambda input: model.apply({'params': params}, input)

 tf_predict = tf.function(

     jax2tf.convert(jax_predict_fn, enable_xla=False),

     input_signature=[

         tf.TensorSpec(

             shape=[1, common.BOARD_SIZE, common.BOARD_SIZE],

             dtype=tf.float32,

             name='input')

     ],

     autograph=False,

 )

 converter = tf.lite.TFLiteConverter.from_concrete_functions(

     [tf_predict.get_concrete_function()], tf_predict)

 tflite_model = converter.convert()

 # Save the model

 with open(os.path.join(modeldir, 'planestrike.tflite'), 'wb') as f:

   f.write(tflite_model)

convertBoardStateToByteBuffer(board);
tflite.run(boardData, outputProbArrays);
float[] probArray = outputProbArrays[0];
int agentStrikePosition = -1;
float maxProb = 0;
for (int i = 0; i < probArray.length; i++) {
  int x = i / Constants.BOARD_SIZE;
  int y = i % Constants.BOARD_SIZE;
  if (board[x][y] == BoardCellStatus.UNTRIED && probArray[i] > maxProb) {
    agentStrikePosition = i;
    maxProb = probArray[i];
  }
}