如何将PPOClipAgent的ActorDistributionNetwork转为输出密度函数的TensorFlow Lite模型？

问题根源 #

省略signature参数保存模型时，TensorFlow会自动生成多个签名，导致TFLite转换时报错Only support a single signature key；同时默认导出的是采样动作而非分布参数，不符合你需要的动作密度函数需求。

解决步骤 #

1. 定义签名函数（返回分布参数）

明确你的动作分布类型（比如连续动作用Normal分布，离散动作用Categorical分布），定义一个只返回分布核心参数的函数，明确输入输出的张量规格：

import tensorflow as tf

# 假设actor_net是训练完成的ActorDistributionNetwork
@tf.function(input_signature=[
    tf.TensorSpec(
        shape=actor_net.input_spec.shape, 
        dtype=actor_net.input_spec.dtype, 
        name="observation"
    )
])
def export_distribution_params(observation):
    # 用推理模式调用网络，得到分布对象
    dist = actor_net(observation, training=False)
    # 根据分布类型返回参数：
    # 连续动作（Normal分布）返回均值和标准差
    return {"mean": dist.mean(), "stddev": dist.stddev()}
    # 离散动作（Categorical分布）返回logits
    # return {"logits": dist.logits()}

2. 带唯一签名保存SavedModel

显式指定serving_default签名，确保只有一个签名被保存：

saved_model_path = "./actor_distribution_net"
tf.saved_model.save(
    actor_net,
    saved_model_path,
    signatures={"serving_default": export_distribution_params}
)

3. 转换为TFLite模型

此时转换不会出现多签名报错，且模型输出的是动作分布的参数：

# 初始化转换器
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path)
# 可选：开启优化（比如量化）
# converter.optimizations = [tf.lite.Optimize.DEFAULT]

# 转换模型
tflite_model = converter.convert()

# 保存TFLite文件
with open("./actor_distribution_net.tflite", "wb") as f:
    f.write(tflite_model)

4. 验证TFLite模型输出

加载TFLite模型后，输入观测即可得到分布参数，你可以基于这些参数计算动作的密度函数（比如Normal分布的概率密度，Categorical分布的概率质量）。

内容的提问来源于stack exchange，提问作者Setjmp