加载TensorFlow策略调用action方法时触发ValueError问题求助

我帮你分析下这个问题的根源，其实是你创建的TimeStep里的部分字段没有匹配策略预期的批量维度要求。

先看错误信息里的核心矛盾：

预期的Option 1和Option 2中，TimeStep的step_type、reward、discount都是shape=(None,)的张量（支持任意批量大小的一维张量），但你传入的这些字段都是标量张量（shape=()），只有observation是shape=(1,4)符合要求，这就导致整体输入和策略保存时的签名不匹配，触发了找不到对应函数的错误。

解决方法：给TimeStep的所有字段添加批量维度 #

你需要把step_type、reward、discount都从标量改成一维张量（比如shape=(1,)），和observation的批量维度保持一致。这里提供两种修改方式：

方式1：手动构造带批量维度的TimeStep

import numpy as np
import tensorflow as tf
from tf_agents.trajectories import time_step as ts
import os
import tempfile

# Temporary directory for the policy
tempdir = os.getenv("TEST_TMPDIR", tempfile.gettempdir())
policy_dir = os.path.join(tempdir, 'policy')

# Load the saved policy
saved_policy = tf.saved_model.load(policy_dir)

# Define state with batch dimension (shape (1, 4))
state = np.array([[0, 0.5, 1, 0]], dtype=np.float32)

# 手动构造所有字段都带批量维度的TimeStep
time_step = ts.TimeStep(
    step_type=tf.convert_to_tensor([ts.StepType.FIRST], dtype=tf.int32),  # shape=(1,)
    reward=tf.convert_to_tensor([0.0], dtype=tf.float32),  # shape=(1,)
    discount=tf.convert_to_tensor([1.0], dtype=tf.float32),  # shape=(1,)
    observation=state  # shape=(1,4)
)

# Get the action from the policy
action = saved_policy.action(time_step)

print(action)

方式2：扩展默认TimeStep的字段维度

如果不想手动构造，也可以基于ts.restart生成的TimeStep，给标量字段添加一个维度：

# ... 前面的加载代码不变 ...
time_step = ts.restart(state)
# 给每个标量字段扩展批量维度
time_step = time_step._replace(
    step_type=tf.expand_dims(time_step.step_type, axis=0),
    reward=tf.expand_dims(time_step.reward, axis=0),
    discount=tf.expand_dims(time_step.discount, axis=0)
)

action = saved_policy.action(time_step)
print(action)

为什么会出现这个问题？ #

你训练并保存策略的时候，策略接收的是批量形式的输入（比如训练时会一次性传入多个样本，所以每个TimeStep字段都带有批量维度），因此保存的策略签名要求输入的每个组件都必须有批量维度。而ts.restart(state)默认生成的step_type、reward、discount都是标量（没有批量维度），这就和策略预期的输入签名不匹配了。

备注：内容来源于stack exchange，提问作者user28859194