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;
}