Status PolymetisControllerServerImpl::UpdateController()

in polymetis/src/polymetis_server.cpp [258:304]


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

  // Read chunks of the binary serialized controller params container.
  updates_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++) {
      updates_model_buffer_.push_back(binary_blob[i]);
    }
  }

  // Load param container
  torch::jit::script::Module param_dict_container;
  memstream model_stream(updates_model_buffer_.data(),
                         updates_model_buffer_.size());
  try {
    param_dict_container = torch::jit::load(model_stream);
  } catch (const c10::Error &e) {
    std::cerr << "error loading the param container:\n";
    std::cerr << e.msg() << std::endl;

    return Status::CANCELLED;
  }

  // Create controller update input dict
  param_dict_input_.clear();
  param_dict_input_.push_back(param_dict_container.forward(empty_input_));

  // Update controller & set intervals
  if (custom_controller_context_.status != UNINITIALIZED) {
    custom_controller_context_.controller_mtx.lock();
    custom_controller_context_.custom_controller.get_method("update")(
        param_dict_input_);
    interval->set_start(robot_state_buffer_.size());
    custom_controller_context_.controller_mtx.unlock();
  } else {
    interval->set_start(-1);
  }

  interval->set_end(-1);

  return Status::OK;
}