feat(backend): PolicyRunner ONNX inference wrapper
This commit is contained in:
52
Backend/PolicyRunner.cs
Normal file
52
Backend/PolicyRunner.cs
Normal file
@@ -0,0 +1,52 @@
|
||||
using Microsoft.ML.OnnxRuntime;
|
||||
using Microsoft.ML.OnnxRuntime.Tensors;
|
||||
|
||||
namespace Backend;
|
||||
|
||||
/// <summary>
|
||||
/// Loads a PPO policy ONNX model once and provides deterministic action selection.
|
||||
/// Thread-safe: <see cref="InferenceSession"/> is safe for concurrent Run() calls.
|
||||
/// </summary>
|
||||
public sealed class PolicyRunner : IDisposable
|
||||
{
|
||||
public const int ObsDim = 7;
|
||||
public const int NActions = 4;
|
||||
|
||||
private readonly InferenceSession _session;
|
||||
private readonly string _inputName;
|
||||
|
||||
public PolicyRunner(string onnxPath)
|
||||
{
|
||||
if (!File.Exists(onnxPath))
|
||||
throw new FileNotFoundException($"ONNX model not found at {onnxPath}");
|
||||
_session = new InferenceSession(onnxPath);
|
||||
_inputName = _session.InputMetadata.Keys.First();
|
||||
}
|
||||
|
||||
/// <summary>Run one inference and return the argmax action index.</summary>
|
||||
public int SelectAction(ReadOnlySpan<float> obs)
|
||||
{
|
||||
if (obs.Length != ObsDim)
|
||||
throw new ArgumentException($"expected obs of length {ObsDim}, got {obs.Length}", nameof(obs));
|
||||
|
||||
var tensor = new DenseTensor<float>(new[] { 1, ObsDim });
|
||||
for (int i = 0; i < ObsDim; i++) tensor[0, i] = obs[i];
|
||||
|
||||
using var results = _session.Run(new[]
|
||||
{
|
||||
NamedOnnxValue.CreateFromTensor(_inputName, tensor)
|
||||
});
|
||||
|
||||
var logits = results.First().AsEnumerable<float>().ToArray();
|
||||
// argmax
|
||||
int best = 0;
|
||||
float bestVal = logits[0];
|
||||
for (int i = 1; i < logits.Length; i++)
|
||||
{
|
||||
if (logits[i] > bestVal) { best = i; bestVal = logits[i]; }
|
||||
}
|
||||
return best;
|
||||
}
|
||||
|
||||
public void Dispose() => _session.Dispose();
|
||||
}
|
||||
Reference in New Issue
Block a user