Running SB3 developed agents on TFLite or Coral
Introduction
I’ve been using Stable-Baselines3
to train agents against some custom Gyms, some of which require fairly
large NNs in order to be effective.
I want those agents to eventually be run on a pi or similar, so I need
to export all the way to TFLite
and ideally a Coral.
How to use
Setup
You will need to have configured the Coral system-wide stuff.
Build a venv:
python3 -m venv venv
source venv/bin/activate
python3 -m pip install -r requirements.txt
Running
This comes with enough defaults to do cradle-to-grave demonstration,
but all the pieces take command-line arguments so I can adjust to taste
for my actual use case.
# Train an agent with SB3
python3 ./train.py
# Convert model
python3 ./model_conv.py
# Run original SB3 model
python3 ./run_sb3.py
# Run the onnx model
python3 ./run_onnx.py
# Run the TFLite model
python3 ./run_tflite.py
# Run the Coral model ["edgetpu" in the name will attempt to load Coral]
python3 ./run_tflite.py MountainCarContinuous-v0 model_quant_edgetpu
Cheers,
Gary [email protected]
