53 lines
1.7 KiB
C#
53 lines
1.7 KiB
C#
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();
|
|
}
|