onnx export: properly handle obs_keys

[Nijco]

At the moment, the onnx export kinda ignores the passed observation keys and uses a hard coded "obs" at export. With these minor changes this is fixed and the onnx model has the properly named inputs set after export.

I viewed the exported model in Netron and ran an exported model in godot with the godot addon (I had to modify the add-on to not take the key "obs" for granted, probably will make a pull request for that later) to verify this modification works as intended.

Pls let me know if any further changes are necessary to get this merged.

Kind regards
Nico

[Ivan-267]

Hello, and thanks for this PR.

A couple of quick test notes:

1. When trying to export multiple keys, e.g. obs1 and obs2 using export_model_as_onnx(model, str(path_onnx), obs_keys=["obs1, obs2"]) I got: KeyError: 'obs1'. This wasn't previously supported either, so it's not something that this PR needs to address.
2. With a single key, the export works, but there's a verification error, as our verification currently uses "obs" for the name.

I should add for context that the original code exported ["obs"] and the recent change to support changing the name was only intended to allow training with a custom obs dictionary name (#212), it did still use ["obs"] for inference, as it didn't require changing the onnx code and the dictionary name during inference shouldn't have any effect. The PR didn't tackle support for multiple keys.

When it comes to multiple obs dictionaries, in the simplest case the observations are just concatenated, so it should be possible to use a single dictionary as well. In other cases where each observation dictionary is processed separately it is different (currently not supported by our onnx inference code, but Python inference support is available).

Other considerations (I edited this section after re-checking the plugin PR and seeing it shouldn't affect compatibility):

• We should also check also if the CleanRL and Rllib exported onnx still work with the plugin changes.

Edit 2: I've briefly checked with clean_rl_pqn_example for now, it seems to work. I didn't check the plugin code in depth yet, but I assume it just sends the observation arrays for each dictionary key to the onnx model instead of the fixed "obs", and so should still work with previously exported onnx files.

[Nijco]

Hello, thank you for locking into my PR.

About your notes:

1. Multiple keys should be supported, you made a slight typo: obs_keys=["obs1, obs2"]), should be obs_keys=["obs1", "obs2"]). Multiple keys in the obs-dict was the motivation to do this, as I use this to have a custom feature extractors for a specific key in the exported ONNX model.
2. Oh yea, when I made this work, I just deactivated the verification and overlooked that when I made this PR. I added a commit to support custom & multiple keys in the verification step.

I think it's quite relevant to support multiple keys for stuff like the aforementioned feature extractor (or at least I wouldn't know how to do this in such a simple manner else).

The exported ONNX models of the other libraries should absolutely still work in the plugin with that PR, as the behavior is still pretty much the same. Exactly like you said, it just does what the plugin did for the 'obs' key before for all the keys in the observation dictionary now.

[Ivan-267]

Thanks for the update and nice catch about the typo.

Everything sounds good based on your comments and I agree with the changes.

Multiple keys in the obs-dict was the motivation to do this, as I use this to have a custom feature extractors for a specific key in the exported ONNX model.

Yes, it's likely the best option for those cases.

Thanks again for the PR. I will start the automated workflows, it may be a couple of days until I can manually check it for a final test, but otherwise it all looks good.

[Ivan-267]

Thanks for clarifying on some points. I performed a quick test, should be ready to merge after adding support for use_obs_array=True as above.

[Ivan-267]

@edbeeching should be ready to merge now.

[edbeeching]

merged, thanks