Stable-baselines: Feature Extraction Causing Problems

[DroneMesh]

Hi Everyone,

So I have noticed an issue when exporting due to the feature extraction. I am guessing that even if you are not using the image recognition part of the policy it still sets up those nodes in the graph. Now the issue arises when you want to export this to an ARM processor. tfcompile will automatically pickup that there is some sort conv network due to the nodes and require the arm compute library even though you do not neet it and thus making the hex much larger for no reason and in most cases the hex will not fit the arm mcu.

Any help into removing these and or way to turn that part off would be great thank you.

[Miffyli]

I am not sure why "feature extraction" would cause issues per-se, but I understand if convolutional layers would do this.

In any case, you might want to take a look at exporting models documentation. I recommend you only extract the network weights with agent.get_parameters(), and on your target machine re-create the network manually and then assign those parameters to correct places.

[DroneMesh]

@Miffyli Thank you for your quick response. I am not experienced enough to recreate the network manually at the moment but is there is a way to remove them. If you can you point me in the correct way and or help me. I would be able to create the docs for exporting into ARM devices and or C/C++

[Miffyli]

Unfortunately I am not familiar working with NNs in such platforms. However, if your network only consists of 1D input, 1D output and fully-connected layers ("MlpPolicy"), then you only need functions for matrix multiplication and any activation functions (tanh, in the case of stable-baselines), as well any processing needed for the output (e.g. softmax for discrete actions). You can google for "Neural network in C++" or "MLP in C++" to find tutorials on this. Note that you do not need the "gradient" or "backpropagation" part (nor the derivatives), you can assume you already have the weights.

[DroneMesh]

@Miffyli I am using PPO2, MLB, tanh, continuous action space. Can you point me to the source where the network gets initialized i was just looking at tensorboard and I think it's the flatten function at is invoking tfcompile to think this is a conv network?

[DroneMesh]

Ok so after digging some more I figured out the issue. There is no way to bypass the conv part of the policy initialization. The issue is arising from CNN and the MLP_Extracter in the policy.py so when tfcompile sees those parts in the graph it assumes its a conv network and asks for the libraries to be put into the compiled hex file and there is no space to add that library.

Anyone know a workaround to not have them initials in the model.

This makes this library unusable in most cases on small STM32 ARM processors which is very bad because this has great potential.

[Miffyli]

Again, I am not familiar with tfcompile and how it works, but you could try modifying the code by removing all references to convolutional networks (i.e. all imports of convolutional operations). If you only use "MlpPolicy", this should be rather easy. Other than that I do not have ideas to share.

[DroneMesh]

@Miffyli Exactly the case can you help me understand how to bypass that part i have been trying for paste 4 hours now. Any suggestions would be very useful and it will allow me to create a PR on exporting model.

[araffin]

I would suggest you to use SB3, the PyTorch version of SB, currrently in beta: https://github.com/DLR-RM/stable-baselines3
I have successfully installed it on a raspberry pi.

You can also compile the C++ lib only and easily export trained network to c++ with pytorch (cf pytorch doc).
In the documentation, we show also how to export Stable-Baselines2 models to pytorch: https://stable-baselines.readthedocs.io/en/master/guide/export.html

In the past, I also wrote some code to use trained cnn model in C++, using eigen or pytorch c api.
However, I was still using pybind11 to pass python objects to the c++ code.

[DroneMesh]

@araffin I would love too however I am constrained with the TF version and AOT tfcompile. If i change my setup it will set me back a couple of months.

[Miffyli]

@DroneMesh Again, I do not know what is required to make this work, but to you should be able to get rid of all CNN stuff by removing references in policies (leave only MlpPolicy and remove all "cnn" stuff/set to None. Also remove nature_cnn) and remove conv stuff in tf_layers. That should be all cnn-related operations. Not sure why these would be included in the compiling if they are not used, but I am not familiar with tfcompile, and do not have time to delve deeper into this right now.

[DroneMesh]

@Miffyli I ended up removing everything related to CNN stuff however it turns out the issue is with the tensor values and nodes. tfcompile thinks there are long long int in the graph so I thought maybe the tensors are np.float64 but that was not the case. So i am digging some more into it now.

[DroneMesh]

Solved it was an LLVM bug