Closed gilzamir18 closed 11 months ago
Requirements for this feature to be considered fully implemented:
A first step has been achieved. I was able to convert a PyTorch model to Onnx. The next step is to implement a decision loop (Controller with LocalBrain) that successfully performs an inference on the C# side in Unity/Godot.
import torch as th
from torch import nn
from stable_baselines3 import SAC
import gym
class OnnxablePolicy(th.nn.Module):
def __init__(self, actor: th.nn.Module, extractors: nn.ModuleDict):
super().__init__()
# Removing the flatten layer because it can't be onnxed
self.extractors = extractors
self.actor = th.nn.Sequential(
actor.latent_pi,
actor.mu,
# For gSDE
# th.nn.Hardtanh(min_val=-actor.clip_mean, max_val=actor.clip_mean),
# Squash the output
th.nn.Tanh(),
)
def forward(self, observations):
input_tensor = self.extractors(observations)
# Return a (B, self._features_dim) PyTorch tensor, where B is batch dimension.
#input_tensor = th.cat(encoded_tensor_list, dim=1)
# NOTE: You may have to process (normalize) observation in the correct
# way before using this. See `common.preprocessing.preprocess_obs`
return self.actor(input_tensor)
# Example: model = SAC("MlpPolicy", "Pendulum-v1")
model = SAC.load("sac_ai4u.zip", device="cpu")
onnxable_model = OnnxablePolicy(model.policy.actor, model.policy.actor.features_extractor)
dummy_inputs = {k: th.randn(1, *obs.shape) for k, obs in model.observation_space.items()}
th.onnx.export(
onnxable_model,
(dummy_inputs, {}),
"my_sac_actor.onnx",
opset_version=9,
input_names=["input"],
)
This reference is need to get success implementing ONNX model inference in Unity: https://github.com/dotnet/machinelearning/issues/1886
There is no proper way to install ML.NET in a Unity project. The solutions for this are very clunky and require workarounds. The best solution for Unity is to use Barracuda. It remains to be investigated whether it is easy to install ML.NET in Godot. In this case, I will create a unique interface (Facade Design Pattern) that isolates AI4U from the specific ONNX runtime used on each platform (Godot and Unity). And, communicating with this interface, I will create a common inference interface between Godot and Unity.
Installing Nuget packages on Godot is a trivial task, natively supported in the Godot 3.5.2 project (Mono version). Using the Microsoft.ML.OnnxRuntime abstractions seems to be enough for us to handle the problem of creating an inference for an offline model. To do so, run the following commands inside the project directory (assuming the .NET CLI is installed):
dotnet add package Microsoft.ML
dotnet add package Microsoft.ML.OnnxRuntime
A sample of a successiful inference is here.
Support of the ONNX model was implemented in current main branch.
Study the implementation of a local controller completely in .Net to run an agent without needing a Python script.
References:
https://github.com/microsoft/onnxruntime-openenclave/blob/openenclave-public/docs/CSharp_API.md https://github.com/DLR-RM/stable-baselines3/issues/383