Missing documentation to export to ONNX

[Ademord]

Is your feature request related to a problem? Please describe. I have been in touch with people at Barracuda https://github.com/Unity-Technologies/barracuda-release/issues/226 and we figured out:

1. there is nothing in the ML documentation related to the hardcoded output names, nor
2. a proper explanation how to export an SB3 trained model from GymUnityWrapper into ONNX to be able to bring it back to unity.

Describe the solution you'd like I am looking at the output names but I figured out I should also come to officially request the export command here since I have passed many posts over the past two months and reached out to the people who posted them and nobody got actual help.

Describe alternatives you've considered Nobody cannot import a torch exported onnx back into unity if you do not provide the proper way to do it.

Appendix

Related Links, mentioned in the issue above

• https://github.com/Unity-Technologies/ml-agents/blob/d34f3cd6ee078782b22341e4ceb958359069ab60/ml-agents/mlagents/trainers/torch/model_serialization.py#L120
• https://github.com/Unity-Technologies/ml-agents/blob/97eb505f9a34597563aadcde9f9e10479de3ac2b/ml-agents/mlagents/trainers/torch/networks.py#L585

Code to Export to an ONNX from torch without the output info (as help)

    @measure
    def _export_pipeline(self, model, env):
        print("\n\n===============================================================================")
        print("Exporting model...")
        print("------------------------------------------------------------------------------")

        # temp
        full_model = self._get_full_model_for_export(model, env)
        s = self._get_empty_observations_for_export(env)

        # export to onnx
        export_path = os.path.join(self.config.run_dir, "model.onnx")
        torch.onnx.export(full_model, s, export_path)
        print("Model successfully exported." if os.path.isfile(export_path) else "Error: Could not export to onnx.")
        print("------------------------------------------------------------------------------")

    def _get_full_model_for_export(self, model, custom_env):
        print("LOG: type of model: ", type(model))
        return FullModel(
            model.policy.features_extractor,
            model.policy.mlp_extractor,
            model.policy.action_net,
            model.policy.value_net,
            [custom_env.observation_space[i].shape for i in custom_env.observation_space.sample().keys()]
        )

    def _get_empty_observations_for_export(self, env):
        # adapt environment and observations
        s = env.reset()
        print("state shapes:", s[0].shape)

        x = torch.randn(1, *env.observation_space[0].shape).cuda()
        y = torch.randn(1, *env.observation_space[1].shape).cuda()
        z = torch.randn(1, *env.observation_space[2].shape).cuda()
        s = [x, y, z]

        print("--------------------------")
        print(s[0].shape)
        print(x.shape)
        print(y.shape)
        print(z.shape)
        print("--------------------------")

        return s

[mmattar]

Thanks for sharing this @Ademord! I've logged your request internally (MLA-2249) so we can prioritize this accordingly. To confirm, are you currently blocked or has your specific issue been resolved?

[Ademord]

@mmattar : Thanks for your reply. Unfortunately, I have been blocked by this only issue for a couple weeks now.

I would highly appreciate it if you could tell me how to adapt this torch-exported ONNX back compatible to Unity. I see there are hardcoded names etc on the network you export, as described above, but I don't know how to add them myself.

[batu]

@Ademord, I am assuming SB3 refers to stable-baselines3. I currently can run my sb3 trained networks in Unity, I might be able to help you out. Here is a relevant conversation on the sb3 repo that might come in handy: https://github.com/DLR-RM/stable-baselines3/issues/383

Here is my code that I use for inference. It is... janky but works.

[Ademord]

Hi @batu I saw the post you reference and the documentation that was updated based on your Onnxable code. Also this code your reference now, is for importing a model for doing inference, but what here is supposed to happen is that MLAgents accepts an externally (SB3, yes stable baselines 3) trained ONNX, as referenced in first post of this thread. It seems like they have hardcoded words into the heads of the networks, so if someone could tell how to export an ONNX properly from torch to be used with MLAgents and not in general in Unity, hence why the Barracuda team moved this issue here and did not keep it there.

[Ademord]

@mmattar anyone who could support on this in the meantime so my block is relieved? I can help you guys with the documentation later when I have it working, but this is essential for my pipeline atm.

[andrewcoh]

Hi @Ademord

Thank you very much for raising this as it is a really important point for us to settle for the future of ML-Agents to be used with external codebases/trainers. This is our model serialization code. It should give you some insight into the expected tensors and possibly can be used directly with some modification in your use case.

My plan is to follow up with the team on how best to handle this moving forward because we would like to provide for this use case. An option may be to provide a script that takes one onnx model and exports another with the proper structure for ML-A, though this may be hard to do in general. Additionally, we also have the constraints of which operators Barracuda supports.

[andrewcoh]

Alternatively, you can try running your model in inference with python just as ML-Agents does with the --inference CLI command. Then, your onnx model would not need to be compatible with the C# model requirements of ML-Agents.

Is this feasible for you or do you need the model to be part of the game build?

[Ademord]

@andrewcoh thanks for the alternatives. I think a script that updates an exported ONNX to match a unity-expected ONNX would be great. The second alternative is probably the easiest where I forget a possible import back to unity and run inference on the env with visuals on.

🟩 I have been integrating mlagents gymwrapped environments into a fully functional wandb workflow with start, resume, reinitialize capabilities, which was a lot of things to take into account.

🟥 Right now I am dealing with an mlagents error mentioned in https://github.com/wandb/client/issues/2750, which is making me lose my mind, but that is a topic outside of this thread. But this is my biggest problem right now when trying to use an MLAgents Environment outside of MLAgents. -> I think it is critical, based on the paper you guys made which is a huge point of my thesis, to use unity and MLAgents to create new environments/benchmarks and then to have a functional workflow when training with something like SB3 and wandb. I have done * all * this functionality and will pass it to you guys as soon as I finish the thesis but at least this current and hopefully last error is the one that I dread every morning when I wake up after leaving it train overnight in my server. I am in parallel training on my windows but that is a lot slower in compared to a server... If you have any hints on why this error continues to show up please let me know 🙏🏻

🟥 But then for now my windows-trained models will be divorced from my linux-trained models (mlagents and SB3, respectively)...

Update: I guess I am going to give up on trying to make this whole unity_gym + SB3 + wandb thing work, and I will continue my training with mlagents on the linux server. The limitations of mlagents I guess do not apply that hard to my scope and I can work around them in the document instead of trying to implement this. Figure below is my tensorboard progress from the past 30min.

I guess if this will be the end of our exchange, would you have some tips on how to train faster? I am already using a timescale of 50 trying to not break the physics engine calculations, and 5 envs (do envs speed up training?, or is it better to go for a multi-agent environment with copy/paste of agents?)