Error while tracing for onnx export. "Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: numpy.float32"

[SunandanAdhikary]

While trying to export the trained agent network in onnx format using inbuilt tracing function torch.export, facing the following error.

import torch.onnx

from ddpg_agent import Agent

from ACVEnv import *

import os
attacker_agent = Agent(agent_name = 'attacker',alpha=0.0001, beta=0.001, 
            input_dims=a_env.observation_space('attacker').shape, tau=0.001,
            batch_size=64, fc1_dims=400, fc2_dims=300, 
            n_actions=a_env.action_space('attacker').shape[0])

detector_agent = Agent(agent_name='detector' ,
                alpha=0.0001, beta=0.001, 
                input_dims=a_env.observation_space('detector').shape, tau=0.001,
                batch_size=64, fc1_dims=400, fc2_dims=300, 
                n_actions=a_env.action_space('detector').shape[0])

# attacker_agent(a_env).load_models()
attacker_agent.load_models()
detector_agent.load_models()

# attacker_agent(a_env).eval()
attacker_agent.eval()
detector_agent.eval()

a_env.reset()

observations = {'attacker':a_env._observation_spaces["attacker"].sample(), 'detector':a_env._observation_spaces["detector"].sample()}

agent_names = ["attacker", "detector"]

agents = {'attacker':attacker_agent, 'detector':detector_agent}

for agent in agent_names:
    
    file_path = os.path.join('tmp/onnx', agent+'_exported.onnx')
    print(torch.flatten(torch.tensor(observations[agent])))
    torch.onnx.export(agents[agent],  tuple(observations[agent]),  file_path,  verbose=True)

The corresponding issue is closed but found no solid solution!

RuntimeError: Only tuples, lists and Variables supported as JIT inputs/outputs. Dictionaries and strings are also accepted but their usage is not recommended. But got unsupported type NoneType

def forward(self, observation):
        self.actor.eval()
        state = T.tensor([observation], dtype = T.float).to(self.actor.device)
        mu = self.actor.forward(state).to(self.actor.device)
        mu_prime = mu + T.tensor(self.noise(), 
                                dtype=T.float).to(self.actor.device)
        self.actor.train()

        return mu_prime

…also adding the forward function defined for the ddpg agrnt if that helps.

[ptrblck]

Based on the error message it seems your model returns a numpy.float32, which is not supported. Convert it to a tensor and this error might be solved.

[SunandanAdhikary]

No, we tried giving observation inputs in tuple/tensor formats, but a similar message appears. The last error stack trace shows it is thrown from somewhere inside jit. We have also raised an issue regarding the same here which contains the complete error log for your reference.

Type error while tracing for exporting pytorch RL agents to ONNX