ONNX support for decision transformers

[ra9hur]

Feature request

I am trying to train off-line RL using decision transformer, convert to .onnx.

from pathlib import Path
from transformers.onnx import FeaturesManager

feature = "sequence-classification"

# load config
model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=feature)
onnx_config = model_onnx_config(model.config)

# export
onnx_inputs, onnx_outputs = transformers.onnx.export(
        #preprocessor=tokenizer,
        model=model,
        config=onnx_config,
        opset=13,
        output=Path("trained_models/DT-model.onnx")
)

Get the below error:

KeyError: "decision-transformer is not supported yet. Only ['albert', 'bart', 'beit', 'bert', 'big-bird', 'bigbird-pegasus', 'blenderbot', 'blenderbot-small', 'bloom', 'camembert', 'clip', 'codegen', 'convbert', 'convnext', 'data2vec-text', 'data2vec-vision', 'deberta', 'deberta-v2', 'deit', 'detr', 'distilbert', 'electra', 'flaubert', 'gpt2', 'gptj', 'gpt-neo', 'groupvit', 'ibert', 'imagegpt', 'layoutlm', 'layoutlmv3', 'levit', 'longt5', 'longformer', 'marian', 'mbart', 'mobilebert', 'mobilenet-v1', 'mobilenet-v2', 'mobilevit', 'mt5', 'm2m-100', 'owlvit', 'perceiver', 'poolformer', 'rembert', 'resnet', 'roberta', 'roformer', 'segformer', 'squeezebert', 'swin', 't5', 'vision-encoder-decoder', 'vit', 'whisper', 'xlm', 'xlm-roberta', 'yolos'] are supported. If you want to support decision-transformer please propose a PR or open up an issue."

Motivation

I would want to use trained models in Godot-RL-Agents. Currently agents are trained using PPO OR imitation learning and bothe support onnx format. Supporting decision transformers could hugely help training models navigating complex scenarios.

Your contribution

I would be interested to raise a PR. But at this time, I have no idea how to go about this. With little bit of guidance, I can try.

[IlyasMoutawwakil]

Hi ! here's an example of a PR adding export support for a new model #2025 and this one is adding inference support for a new task #2031