Running SB3 developed agents on TFLite or Coral


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


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


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 ./

# Convert model
python3 ./

# Run original SB3 model
python3 ./
# Run the onnx model
python3 ./
# Run the TFLite model
python3 ./
# Run the Coral model ["edgetpu" in the name will attempt to load Coral]
python3 ./ MountainCarContinuous-v0 model_quant_edgetpu

Gary [email protected]


View Github