使用JAX构建强化学习agent并借助TensorFlowLite将其部署到Android应用中

在之前发布文章《一个新 TensorFlow Lite 示例应用：棋盘游戏》中，展示了如何使用 TensorFlow 和 TensorFlow Agents 来训练强化学习 (RL) agent，使其玩一个简单棋盘游戏 “Plane Strike”。我们还将训练后模型转换为 TensorFlow Lite，然后将其部署到功能完备 Android 应用中。本文，我们将演示一种全新路径: 使用 Flax/JAX 训练相同强化学习 agent，然后将其部署到我们之前构建同一款 Android 应用中。我们已经在 tensorflow/examples 代码库中开放了完整源代码以供您参考。

• Flaxhttps://flax.readthedocs.io/
• JAXhttps://jax.readthedocs.io/
• tensorflow/exampleshttps://github.com/tensorflow/examples/blob/master/lite/examples/reinforcement_learning/ml/tf_and_jax/training_jax.py

简单回顾一下游戏规则: 我们基于强化学习 agent 需要根据真人玩家棋盘位置预测击打位置，以便能早于真人玩家完成游戏。如需进一步了解游戏规则，请参阅我们之前发布文章。

△ “Plane Strike” 游戏演示 背景: JAX 和 TensorFlow

JAX 是一个与 NumPy 类似内容库，由 Google Research 部门专为实现高性能计算而开发。JAX 使用 XLA 针对 GPU 和 TPU 优化程序进行编译。

• JAXhttps://github.com/google/jax
• XLAhttps://tensorflow.google.cn/xla
• TPUhttps://cloud.google.com/tpu

而 Flax 则是在 JAX 基础上构建一款热门神经网络库。研究人员一直在使用 JAX/Flax 来训练包含数亿万个参数超大模型 (如用于语言理解和生成 PaLM，或者用于图像生成 Imagen)，以便充分利用现代硬件。

• Flaxhttps://github.com/google/flax
• PaLMhttps://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html
• Imagenhttps://imagen.research.google/

如果您不熟悉 JAX 和 Flax，可以先从 JAX 101 教程和 Flax 入门示例开始。

• JAX 101 教程https://jax.readthedocs.io/en/latest/jax-101/index.html
• Flax 入门示例https://flax.readthedocs.io/en/latest/getting_started.html

2015 年底，TensorFlow 作为 Machine Learning (ML) 内容库问世，现已发展为一个丰富生态系统，其中包含用于实现 ML 流水线生产化 (TFX)、数据可视化 (TensorBoard)，和将 ML 模型部署到边缘设备 (TensorFlow Lite) 工具，以及在网络浏览器上运行装置，或能够执行 JavaScript (TensorFlow.js) 任何装置。

• TFXhttps://tensorflow.google.cn/tfx
• TensorBoardhttps://tensorboard.dev/
• TensorFlow Litehttps://tensorflow.google.cn/lite
• TensorFlow.jshttps://tensorflow.google.cn/js

在 JAX 或 Flax 中开发模型也可以利用这一丰富生态系统。方法是首先将此类模型转换为 TensorFlow SavedModel 格式，然后使用与它们在 TensorFlow 中原生开发相同工具。

• SavedModelhttps://tensorflow.google.cn/guide/saved_model

如果您已经拥有经 JAX 训练模型并希望立即进行部署，我们整合了一份资源列表供您参考:

• 频 “使用 TensorFlow Serving 为 JAX 模型提供服务”，展示了如何使用 TensorFlow Serving 部署 JAX 模型:
• 文章《借助 TensorFlow.js 在网络上使用 JAX》，对如何将 JAX 模型转换为 TFJS，并在网络应用中运行进行了详细讲解: 

  https://blog.tensorflow.org/2022/08/jax-on-web-with-tensorflowjs.html

• 本篇文章演示了如何将 Flax/JAX 模型转换为 TFLite，并在原生 Android 应用中运行该模型。

总而言之，无论您部署目标是服务器、网络还是移动设备，我们都会为您提供相应帮助。使用 Flax/JAX 实现游戏 agent

将目光转回到棋盘游戏。为了实现强化学习 agent，我们将会利用与之前相同 OpenAI gym 环境。这次，我们将使用 Flax/JAX 训练相同策略梯度模型。回想一下，在数学层面上策略梯度定义是: 

• OpenAI gymhttps://github.com/tensorflow/examples/tree/master/lite/examples/reinforcement_learning/ml/tf_and_jax/gym_planestrike/gym_planestrike/envs

其中:

• T: 每段时步数，各段时步数可能有所不同
• st: 时步上状态 t
• at: 时步上所选操作 t 指定状态 s
• πθ: 参数为 θ 策略
• R(*): 在指定策略下，收集到奖励

我们定义了一个 3 层 MLP 作为策略网络，该网络可以预测 agent 下一个击打位置。 class PolicyGradient(nn.Module):

“””Neural network to predict the next strike position.”””

@nn.compact
def __call__(self, x):
dtype = jnp.float32
x = x.reshape((x.shape[0], -1))
x = nn.Dense(
features=2 * common.BOARD_SIZE**2, name=’hidden1′, dtype=dtype)(
x)
x = nn.relu(x)
x = nn.Dense(features=common.BOARD_SIZE**2, name=’hidden2′, dtype=dtype)(x)
x = nn.relu(x)
x = nn.Dense(features=common.BOARD_SIZE**2, name=’logits’, dtype=dtype)(x)
policy_probabilities = nn.softmax(x)
return policy_probabilities

在我们训练循环每次迭代中，我们都会使用神经网络玩一局游戏、收集轨迹信息 (游戏棋盘位置、采取操作和奖励)、对奖励进行折扣，然后使用相应轨迹训练模型。for i in tqdm(range(iterations)):
predict_fn = functools.partial(run_inference, params)
board_log, action_log, result_log = common.play_game(predict_fn)
rewards = common.compute_rewards(result_log)
optimizer, params, opt_state = train_step(optimizer, params, opt_state,
board_log, action_log, rewards)在 train_step() 方法中，我们首先会使用轨迹计算损失，然后使用 jax.grad() 计算梯度，最后，使用 Optax (用于 JAX 梯度处理和优化库) 来更新模型参数。def compute_loss(logits, labels, rewards):
one_hot_labels = jax.nn.one_hot(labels, num_classes=common.BOARD_SIZE**2)
loss = -jnp.mean(
jnp.sum(one_hot_labels * jnp.log(logits), axis=-1) * jnp.asarray(rewards))
return loss

def train_step(model_optimizer, params, opt_state, game_board_log,
predicted_action_log, action_result_log):
“””Run one training step.”””

def loss_fn(model_params):
logits = run_inference(model_params, game_board_log)
loss = compute_loss(logits, predicted_action_log, action_result_log)
return loss

def compute_grads(params):
return jax.grad(loss_fn)(params)

grads = compute_grads(params)
updates, opt_state = model_optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)
return model_optimizer, params, opt_state

@jax.jit
def run_inference(model_params, board):
logits = PolicyGradient().apply({‘params’: model_params}, board)
return logits

• Optaxhttps://github.com/deepmind/optax

这就是训练循环。如下图所示，我们可以在 TensorBoard 中观察训练进度；其中，我们使代理指标 “game_length” (完成游戏所需步骤数) 来跟踪进度: 若 agent 变得更聪明，它便能以更少步骤完成游戏。将 Flax/JAX 模型转换为 TensorFlow Lite 并与 Android 应用集成 
完成模型训练后，我们使用 jax2tf (一款 TensorFlow-JAX 互操作工具)，将 JAX 模型转换为
TensorFlow concrete function。最后一步是调用 TensorFlow Lite 转换器来将 concrete
function 转换为 TFLite 模型。

# Convert to tflite model
model = PolicyGradient()
jax_predict_fn = lambda input: model.apply({‘params’: params}, input)
if_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)

• jax2tf

https://github.com/google/jax/tree/main/jax/experimental/jax2tf

经 JAX 转换 TFLite 模型与任何经 TensorFlow 训练 TFLite 模型会有完全一致行为。您可以使用 Netron 进行可视化:

使用 Netron 对 Flax/JAX 转换 TFLite 模型进行可视化

我们可以使用与之前完全一样 Java 代码来调用模型并获取预测结果。

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];
}
}

总结

本文详细介绍了如何使用 Flax/JAX 训练简单强化学习模型、利用 jax2tf 将其转换为 TensorFlow Lite，以及将转换后模型集成到 Android 应用。

现在，您已经了解了如何使用 Flax/JAX 构建神经网络模型，以及如何利用强大 TensorFlow 生态系统，在几乎任何您想要位置部署模型。我们十分期待看到您使用 JAX 和 TensorFlow 构建出色应用！