Export trained ML-Agents models to ONNX format for Unity deployment and push to HuggingFace Hub for sharing
Export and deploy trained ML-Agents models.
Models are automatically exported during training to ONNX format:
results/<run-id>/
āāā <BehaviorName>.onnx # Exported ONNX model
āāā <BehaviorName>/
ā āāā checkpoint.pt # PyTorch checkpoint
āāā configuration.yaml # Training config
If you need to re-export:
from mlagents.trainers.torch_entities.model_serialization import export_policy_model
# Load checkpoint and export
export_policy_model(
checkpoint_path="results/MyRun/MyBehavior/checkpoint.pt",
output_filepath="results/MyRun/MyBehavior.onnx"
)
# Copy model to Unity Assets
cp results/MyRun/MyBehavior.onnx Project/Assets/ML-Agents/Models/
.onnx file// In Unity, check model is loaded:
var model = GetComponent<BehaviorParameters>().Model;
if (model != null)
{
Debug.Log("Model loaded successfully!");
}
Share your trained model on HuggingFace:
# Set HuggingFace token
export HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxx
# Push model
mlagents-push-to-hf \
--run-id=MyTraining \
--local-dir=results/MyTraining \
--repo-id=username/my-agent-model \
--commit-message="Trained PPO agent on CustomEnv"
The push command automatically generates a model card with:
Download and use community models:
# Download model
mlagents-load-from-hf \
--repo-id=username/my-agent-model \
--local-dir=./downloaded_models
# Copy to Unity
cp downloaded_models/*.onnx Project/Assets/ML-Agents/Models/
Verify exported model works correctly:
import onnx
# Load ONNX model
model = onnx.load("results/MyRun/MyBehavior.onnx")
# Check the model
onnx.checker.check_model(model)
print("Model is valid!")
# Print model info
print(f"Inputs: {[input.name for input in model.graph.input]}")
print(f"Outputs: {[output.name for output in model.graph.output]}")
Reduce model size for deployment:
# In training config, reduce network size:
network_settings:
hidden_units: 64 # Down from 128
num_layers: 2 # Down from 3
Smaller networks:
ModuleNotFoundError: No module named 'onnxscript'
Solution:
# Ensure correct torch version
pip install torch<=2.8.0
results/
āāā MyRun/
āāā MyBehavior.onnx # ā Deploy this to Unity
āāā MyBehavior/
ā āāā checkpoint.pt # ā Keep for resuming training
āāā configuration.yaml # ā Training config reference
āāā events.out.tfevents.* # ā TensorBoard logs
āāā run_logs/
āāā training_status.json # ā Training metadata
Browse ML-Agents models:
Share your models:
mlagents-push-to-hftrain-ml-agent - Train models before exportdebug-training - Fix issues before exportoptimize-performance - Optimize model size/speed