Unity ML-Agents导入自定义ONNX模型报'vector_observation'错误求助

问题 #

在Unity ML-Agents项目中导入自定义PyTorch神经网络模型，导出并修改ONNX后运行Unity报错：

UnityAgentsException: Unknown tensorProxy expected as input : vector_observation
Unity.MLAgents.Inference.TensorGenerator.GenerateTensors (System.Collections.Generic.IReadOnlyList`1[T] tensors, System.Int32 currentBatchSize, System.Collections.Generic.IList`1[T] infos) (at ./Library/PackageCache/com.unity.ml-agents@89a6357016/Runtime/Inference/TensorGenerator.cs:178)
Unity.MLAgents.Inference.ModelRunner.DecideBatch () (at ./Library/PackageCache/com.unity.ml-agents@89a6357016/Runtime/Inference/ModelRunner.cs:210)
Unity.MLAgents.Policies.BarracudaPolicy.DecideAction () (at ./Library/PackageCache/com.unity.ml-agents@89a6357016/Runtime/Policies/BarracudaPolicy.cs:134)
Unity.MLAgents.Agent.DecideAction () (at ./Library/PackageCache/com.unity.ml-agents@89a6357016/Runtime/Agent.cs:1402)
Unity.MLAgents.Academy.EnvironmentStep () (at ./Library/PackageCache/com.unity.ml-agents@89a6357016/Runtime/Academy.cs:597)
Unity.MLAgents.AcademyFixedUpdateStepper.FixedUpdate () (at ./Library/PackageCache/com.unity.ml-agents@89a6357016/Runtime/Academy.cs:43)

自定义模型导出ONNX的Python代码如下：

import torch
import torch.nn as nn
import onnx

device = "cuda" if torch.cuda.is_available() else "cpu"


class EnhancedMLP(nn.Module):
    def __init__(self, input_size, output_size):
        super(EnhancedMLP, self).__init__()

        # Doubling the number of layers and units
        self.fc1 = nn.Linear(input_size, 256)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 256)
        self.fc4 = nn.Linear(256, 256)
        self.fc5 = nn.Linear(256, output_size)

        self.activation = nn.ReLU()

    def forward(self, x):
        output = self.fc5(self.activation(
            self.fc4(self.activation(self.fc3(self.activation(self.fc2(self.activation(self.fc1(x)))))))))
        continuous_action_shape = torch.tensor([2], dtype=torch.int64)  # Assuming 2 continuous actions
        return output, continuous_action_shape


model = EnhancedMLP(input_size=96, output_size=2).to(device)

# set the model to inference mode
model.eval()

# Create a dummy input tensor
dummy_input = torch.randn(1, 96).to(device)

# ... [rest of your code]

# Export the model to ONNX format
torch.onnx.export(model,                       # model being run
                  dummy_input,                 # model input (or a tuple for multiple inputs)
                  "Albert.onnx",               # where to save the model (can be a file or file-like object)
                  export_params=True,          # store the trained parameter weights inside the model file
                  opset_version=9,             # the ONNX version to export the model to
                  do_constant_folding=True,    # whether to execute constant folding for optimization
                  input_names=['vector_observation'],       # the model's input name
                  output_names=['continuous_actions', 'continuous_action_output_shape'],      # the model's output names
                  dynamic_axes={'vector_observation': {0: 'batch_size'},  # dynamic axis for the input tensor
                                'continuous_actions': {0: 'batch_size'},
                                'continuous_action_output_shape': {0: 'batch_size'}})


# Load the exported ONNX model
torchmodel = onnx.load('Albert.onnx')
graph = torchmodel.graph

# Add the version_number field
version_number = onnx.helper.make_tensor("version_number", onnx.TensorProto.INT64, [1], [3])
graph.initializer.append(version_number)

# Add the version_number to the model's outputs
version_number_info = onnx.helper.make_tensor_value_info("version_number", onnx.TensorProto.INT64, shape=[])
graph.output.append(version_number_info)

# Add the memory_size field
memory_size = onnx.helper.make_tensor("memory_size", onnx.TensorProto.INT64, [1], [0])
graph.initializer.append(memory_size)

# Add the memory_size to the model's outputs
memory_size_info = onnx.helper.make_tensor_value_info("memory_size", onnx.TensorProto.INT64, shape=[])
graph.output.append(memory_size_info)

# Define the continuous_action_output_shape tensor value info
continuous_action_shape_info = onnx.helper.make_tensor_value_info(
    "continuous_action_output_shape",
    onnx.TensorProto.INT64,
    [1]
)

# Add the continuous_action_output_shape to the model's outputs
graph.output.append(continuous_action_shape_info)

# Define and add the actual tensor value for continuous_action_output_shape
continuous_action_shape_tensor = onnx.helper.make_tensor(
    "continuous_action_output_shape",
    onnx.TensorProto.INT64,
    [1],
    [2]  # Assuming 2 continuous actions
)
graph.initializer.append(continuous_action_shape_tensor)

# Save the modified ONNX model
onnx.save(torchmodel, 'ModifiedAlbert.onnx')

print("Model has been converted to ONNX and modified.")

# model.eval()
# dummy_input = torch.randn(1, 96)
# dummy_input = dummy_input.to(device)
# torch.onnx.export(model, dummy_input, "AlbertModelMLP01mod.onnx")

# return torch.tensor([1.0, 0.0]).to(device)

model = onnx.load("Albert.onnx")
print(model.graph.input)  # Check the input tensor names
print(model.graph.output) # Check the output tensor names

问题原因分析 #

1. 输入名称未被正确保留：PyTorch导出ONNX时可能自动修改输入张量名称，导致Unity期望的vector_observation不存在。
2. 输出结构重复定义：代码中先在导出参数里声明continuous_action_output_shape为输出，后续又手动添加该张量的初始化和输出信息，导致ONNX模型结构混乱，间接影响输入识别。
3. 低版本ONNX兼容性问题：使用opset_version=9可能与ML-Agents的Barracuda推理引擎存在兼容性问题，无法正确解析输入张量信息。

解决方案 #

1. 确认并修正输入张量名称 #

导出初始ONNX后，通过print(model.graph.input)检查实际输入名称。如果不是vector_observation，加载模型后手动修改：

torchmodel = onnx.load('Albert.onnx')
graph = torchmodel.graph
# 修改第一个输入张量的名称为vector_observation
graph.input[0].name = "vector_observation"

2. 简化ONNX输出修改逻辑 #

移除重复添加continuous_action_output_shape的代码，仅保留version_number和memory_size的添加逻辑：

torchmodel = onnx.load('Albert.onnx')
graph = torchmodel.graph

# 添加version_number
version_number = onnx.helper.make_tensor("version_number", onnx.TensorProto.INT64, [1], [3])
graph.initializer.append(version_number)
version_number_info = onnx.helper.make_tensor_value_info("version_number", onnx.TensorProto.INT64, shape=[])
graph.output.append(version_number_info)

# 添加memory_size
memory_size = onnx.helper.make_tensor("memory_size", onnx.TensorProto.INT64, [1], [0])
graph.initializer.append(memory_size)
memory_size_info = onnx.helper.make_tensor_value_info("memory_size", onnx.TensorProto.INT64, shape=[])
graph.output.append(memory_size_info)

# 保存修改后的模型
onnx.save(torchmodel, 'ModifiedAlbert.onnx')

3. 调整ONNX导出参数 #

提升opset版本到11+，关闭常量折叠避免张量结构被自动修改：

torch.onnx.export(model,
                  dummy_input,
                  "Albert.onnx",
                  export_params=True,
                  opset_version=11,
                  do_constant_folding=False,
                  input_names=['vector_observation'],
                  output_names=['continuous_actions', 'continuous_action_output_shape'],
                  dynamic_axes={'vector_observation': {0: 'batch_size'},
                                'continuous_actions': {0: 'batch_size'},
                                'continuous_action_output_shape': {0: 'batch_size'}})

4. 验证ONNX模型合法性 #

使用ONNX工具检查模型结构是否正确：

import onnx
torchmodel = onnx.load('ModifiedAlbert.onnx')
onnx.checker.check_model(torchmodel)

也可以用Netron可视化工具确认输入名称为vector_observation，输出结构符合要求。

内容的提问来源于stack exchange，提问作者smn.ksg