Status PolymetisControllerServerImpl::SetController()

in polymetis/src/polymetis_server.cpp [213:256]


Status PolymetisControllerServerImpl::SetController(
    ServerContext *context, ServerReader<ControllerChunk> *stream,
    LogInterval *interval) {
  std::lock_guard<std::mutex> service_lock(service_mtx_);

  resetControllerContext();
  custom_controller_context_.server_context = context;

  // Read chunks of the binary serialized controller. The binary messages
  // would be written into the preallocated buffer used for the Torch
  // controllers.
  controller_model_buffer_.clear();
  ControllerChunk chunk;
  while (stream->Read(&chunk)) {
    std::string binary_blob = chunk.torchscript_binary_chunk();
    for (int i = 0; i < binary_blob.size(); i++) {
      controller_model_buffer_.push_back(binary_blob[i]);
    }
  }

  memstream model_stream(controller_model_buffer_.data(),
                         controller_model_buffer_.size());
  try {
    custom_controller_context_.custom_controller =
        torch::jit::load(model_stream);
  } catch (const c10::Error &e) {
    std::cerr << "error loading the model:\n";
    std::cerr << e.msg() << std::endl;

    return Status::CANCELLED;
  }
  custom_controller_context_.status = READY;
  std::cout << "Loaded new controller.\n";

  // Respond with start index
  while (custom_controller_context_.status == READY) {
    usleep(SPIN_INTERVAL_USEC);
  }
  interval->set_start(custom_controller_context_.episode_begin);
  interval->set_end(-1);

  // Return success.
  return Status::OK;
}