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