agent.cpp

Go to the documentation of this file.

//===- agent.cpp - RL Agent/Model for ONNX Runner --------------*- C++ -*-===//
//
// Part of the MLCompilerBridge Project, under the Apache License v2.0 with LLVM
// Exceptions. See the LICENSE file for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#include "MLModelRunner/ONNXModelRunner/agent.h"
#include "MLModelRunner/Utils/Debug.h"
#include <algorithm>
#include <cmath>
#include <iterator>
 
namespace MLBridge {
 
Agent::Agent(std::string modelPath) {
  // Create model object here
  this->model = new ONNXModel(modelPath.c_str());
}
 
unsigned Agent::computeAction(Observation &input) {
  // Call model on input
  assert(input.size() > 0);
  llvm::SmallVector<float, 100> model_input(input.begin(), input.end());
  llvm::SmallVector<float, 100> model_output;
 
  this->model->run(model_input, model_output);
 
  // select action from model output
  auto max = std::max_element(model_output.begin(),
                              model_output.end()); // [2, 4)
  int argmaxVal = std::distance(model_output.begin(), max);
 
  MLBRIDGE_DEBUG(
      std::cout << "---------------MODEL OUTPUT VECTOR:----------------\n";
      for (auto e
           : model_output) { std::cout << e << " "; } std::cout
      << "\nmax value and index are " << *max << " " << argmaxVal << "\n";);
  return argmaxVal;
}
 
} // namespace MLBridge