From 18aaf120ab851ffa115113209337ea92a9484933 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Tue, 15 Aug 2023 09:37:51 -0600 Subject: [PATCH 01/11] add kalman filter, modifications to batch estimator --- CMakeLists.txt | 6 +- grpc/CMakeLists.txt | 57 +- grpc/kalman.proto | 122 ++++ grpc/kalman_server.cc | 58 ++ grpc/kalman_service.cc | 312 +++++++++ grpc/kalman_service.h | 90 +++ mjpc/CMakeLists.txt | 6 +- mjpc/estimators/batch.cc | 593 +++++++++++++----- mjpc/estimators/batch.h | 74 ++- mjpc/estimators/estimator.h | 215 ++++++- mjpc/estimators/gui.cc | 44 ++ mjpc/estimators/gui.h | 54 ++ mjpc/estimators/include.cc | 4 +- mjpc/estimators/include.h | 1 + mjpc/estimators/kalman.cc | 571 +++++++++++++++++ mjpc/estimators/kalman.h | 188 ++++++ mjpc/test/CMakeLists.txt | 16 +- mjpc/test/estimator/CMakeLists.txt | 3 + mjpc/test/estimator/batch_prior_test.cc | 6 +- mjpc/test/estimator/batch_sensor_test.cc | 14 +- .../estimator/estimator_utilities_test.cc | 2 +- mjpc/test/estimator/kalman_test.cc | 115 ++++ mjpc/test/testdata/estimator/box/task2D.xml | 8 + .../test/testdata/estimator/box/task3Drot.xml | 65 ++ .../testdata/estimator/box/task3Drot2.xml | 75 +++ .../testdata/estimator/particle/task_imu.xml | 56 ++ .../testdata/estimator/quadruped/task.xml | 5 - python/mujoco_mpc/batch.py | 2 +- python/mujoco_mpc/kalman.py | 270 ++++++++ python/mujoco_mpc/kalman_test.py | 100 +++ python/setup_kalman.py | 299 +++++++++ 31 files changed, 3236 insertions(+), 195 deletions(-) create mode 100644 grpc/kalman.proto create mode 100644 grpc/kalman_server.cc create mode 100644 grpc/kalman_service.cc create mode 100644 grpc/kalman_service.h create mode 100644 mjpc/estimators/gui.cc create mode 100644 mjpc/estimators/gui.h create mode 100644 mjpc/estimators/kalman.cc create mode 100644 mjpc/estimators/kalman.h create mode 100644 mjpc/test/estimator/kalman_test.cc create mode 100644 mjpc/test/testdata/estimator/box/task3Drot.xml create mode 100644 mjpc/test/testdata/estimator/box/task3Drot2.xml create mode 100644 mjpc/test/testdata/estimator/particle/task_imu.xml create mode 100644 python/mujoco_mpc/kalman.py create mode 100644 python/mujoco_mpc/kalman_test.py create mode 100644 python/setup_kalman.py diff --git a/CMakeLists.txt b/CMakeLists.txt index b00bf380b..90b72be9f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -191,6 +191,6 @@ set(MJPC_COMPILE_OPTIONS "${AVX_COMPILE_OPTIONS}" "${EXTRA_COMPILE_OPTIONS}") set(MJPC_LINK_OPTIONS "${EXTRA_LINK_OPTIONS}") add_subdirectory(mjpc) -if(MJPC_BUILD_GRPC_SERVICE) - add_subdirectory(grpc) -endif() +# if(MJPC_BUILD_GRPC_SERVICE) +# add_subdirectory(grpc) +# endif() diff --git a/grpc/CMakeLists.txt b/grpc/CMakeLists.txt index 37161c3e0..c0cbf465d 100644 --- a/grpc/CMakeLists.txt +++ b/grpc/CMakeLists.txt @@ -55,7 +55,8 @@ get_filename_component(agent_service_proto "./agent.proto" ABSOLUTE) get_filename_component(agent_service_proto_path "${agent_service_proto}" PATH) get_filename_component(batch_service_proto "./batch.proto" ABSOLUTE) get_filename_component(batch_service_proto_path "${batch_service_proto}" PATH) - +get_filename_component(kalman_service_proto "./kalman.proto" ABSOLUTE) +get_filename_component(kalman_service_proto_path "${kalman_service_proto}" PATH) # Generated sources set(agent_service_proto_srcs "${CMAKE_CURRENT_BINARY_DIR}/agent.pb.cc") @@ -68,6 +69,11 @@ set(batch_service_proto_hdrs "${CMAKE_CURRENT_BINARY_DIR}/batch.pb.h") set(batch_service_grpc_srcs "${CMAKE_CURRENT_BINARY_DIR}/batch.grpc.pb.cc") set(batch_service_grpc_hdrs "${CMAKE_CURRENT_BINARY_DIR}/batch.grpc.pb.h") +set(kalman_service_proto_srcs "${CMAKE_CURRENT_BINARY_DIR}/kalman.pb.cc") +set(kalman_service_proto_hdrs "${CMAKE_CURRENT_BINARY_DIR}/kalman.pb.h") +set(kalman_service_grpc_srcs "${CMAKE_CURRENT_BINARY_DIR}/kalman.grpc.pb.cc") +set(kalman_service_grpc_hdrs "${CMAKE_CURRENT_BINARY_DIR}/kalman.grpc.pb.h") + message("We need the following for agent/batch_protos:") message(_GRPC_CPP_PLUGIN_EXECUTABLE${_GRPC_CPP_PLUGIN_EXECUTABLE}) @@ -105,6 +111,23 @@ add_custom_command( "${batch_service_proto}" ) +add_custom_command( + OUTPUT + "${kalman_service_proto_srcs}" + "${kalman_service_proto_hdrs}" + "${kalman_service_grpc_srcs}" + "${kalman_service_grpc_hdrs}" + COMMAND ${_PROTOBUF_PROTOC} + ARGS + --grpc_out "${CMAKE_CURRENT_BINARY_DIR}" + --cpp_out "${CMAKE_CURRENT_BINARY_DIR}" + -I "${kalman_service_proto_path}" + --plugin=protoc-gen-grpc="${_GRPC_CPP_PLUGIN_EXECUTABLE}" + "${kalman_service_proto}" + DEPENDS + "${kalman_service_proto}" +) + # Include generated *.pb.h files include_directories(${CMAKE_CURRENT_BINARY_DIR}/..) @@ -210,3 +233,35 @@ message(BATCH_SERVICE_COMPILE_OPTIONS=${BATCH_SERVICE_COMPILE_OPTIONS}) target_compile_options(batch_server PUBLIC ${BATCH_SERVICE_COMPILE_OPTIONS}) target_link_options(batch_server PRIVATE ${BATCH_SERVICE_LINK_OPTIONS}) +add_executable( + kalman_server + kalman_server.cc + ${kalman_service_grpc_srcs} + ${kalman_service_grpc_hdrs} + ${kalman_service_proto_srcs} + ${kalman_service_proto_hdrs} + kalman_service.h + kalman_service.cc +) + +target_link_libraries( + kalman_server + # kalman_service_grpc_proto + absl::check + absl::flags + absl::flags_parse + absl::log + absl::random_random + absl::status + absl::strings + ${_REFLECTION} + ${_GRPC_GRPCPP} + ${_PROTOBUF_LIBPROTOBUF} + mujoco::mujoco + libmjpc +) + +target_include_directories(kalman_server PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..) +message(KALMAN_SERVICE_COMPILE_OPTIONS=${KALMAN_SERVICE_COMPILE_OPTIONS}) +target_compile_options(kalman_server PUBLIC ${KALMAN_SERVICE_COMPILE_OPTIONS}) +target_link_options(kalman_server PRIVATE ${KALMAN_SERVICE_LINK_OPTIONS}) \ No newline at end of file diff --git a/grpc/kalman.proto b/grpc/kalman.proto new file mode 100644 index 000000000..3e6e24bb2 --- /dev/null +++ b/grpc/kalman.proto @@ -0,0 +1,122 @@ +// Copyright 2023 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +syntax = "proto3"; + +package kalman; + +service Kalman { + // Initialize Kalman + rpc Init(InitRequest) returns (InitResponse); + // Reset Kalman + rpc Reset(ResetRequest) returns (ResetResponse); + // Kalman settings + rpc Settings(SettingsRequest) returns (SettingsResponse); + // Kalman measurement update + rpc UpdateMeasurement(UpdateMeasurementRequest) returns (UpdateMeasurementResponse); + // Kalman prediction update + rpc UpdatePrediction(UpdatePredictionRequest) returns (UpdatePredictionResponse); + // Kalman timers + rpc Timers(TimersRequest) returns (TimersResponse); + // Kalman state + rpc State(StateRequest) returns (StateResponse); + // Kalman covariance + rpc Covariance(CovarianceRequest) returns (CovarianceResponse); + // Kalman noise + rpc Noise(NoiseRequest) returns (NoiseResponse); +} + +message MjModel { + optional bytes mjb = 1; + optional string xml = 2; +} + +message InitRequest { + optional MjModel model = 1; +} + +message InitResponse {} + +message Settings { + optional double epsilon = 1; + optional bool flg_centered = 2; +} + +message SettingsRequest { + optional Settings settings = 1; +} + +message SettingsResponse { + Settings settings = 1; +} + +message ResetRequest {} + +message ResetResponse {} + +message UpdateMeasurementRequest { + repeated double ctrl = 1 [packed = true]; + repeated double sensor = 2 [packed = true]; +} + +message UpdateMeasurementResponse {} + +message UpdatePredictionRequest {} + +message UpdatePredictionResponse {} + +message TimersRequest {} + +message TimersResponse { + double measurement = 1; + double prediction = 2; +} + +message State { + repeated double state = 1 [packed = true]; +} + +message StateRequest { + optional State state = 1; +} + +message StateResponse { + State state = 1; +} + +message Covariance { + repeated double covariance = 1 [packed = true]; + optional int32 dimension = 2; +} + +message CovarianceRequest { + optional Covariance covariance = 1; +} + +message CovarianceResponse { + Covariance covariance = 1; +} + +message Noise { + repeated double process = 1 [packed = true]; + repeated double sensor = 2 [packed = true]; +} + +message NoiseRequest { + optional Noise noise = 1; +} + +message NoiseResponse { + Noise noise = 2; +} \ No newline at end of file diff --git a/grpc/kalman_server.cc b/grpc/kalman_server.cc new file mode 100644 index 000000000..008bba5f9 --- /dev/null +++ b/grpc/kalman_server.cc @@ -0,0 +1,58 @@ +// Copyright 2023 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Startup code for `Kalman` server. + +#include +#include +#include + +#include +#include +#include +#include +#include + +// DEEPMIND INTERNAL IMPORT +#include +#include +#include +#include "grpc/kalman_service.h" + +ABSL_FLAG(int32_t, mjpc_port, 10000, "port to listen on"); + +int main(int argc, char** argv) { + absl::ParseCommandLine(argc, argv); + absl::ParseCommandLine(argc, argv); + int port = absl::GetFlag(FLAGS_mjpc_port); + + std::string server_address = absl::StrCat("[::]:", port); + + std::shared_ptr server_credentials = + grpc::experimental::LocalServerCredentials(LOCAL_TCP); + grpc::ServerBuilder builder; + builder.AddListeningPort(server_address, server_credentials); + + kalman_grpc::KalmanService service; + builder.SetMaxReceiveMessageSize(40 * 1024 * 1024); + builder.RegisterService(&service); + + std::unique_ptr server(builder.BuildAndStart()); + LOG(INFO) << "Server listening on " << server_address; + + // Keep the program running until the server shuts down. + server->Wait(); + + return 0; +} \ No newline at end of file diff --git a/grpc/kalman_service.cc b/grpc/kalman_service.cc new file mode 100644 index 000000000..e8da2f3e2 --- /dev/null +++ b/grpc/kalman_service.cc @@ -0,0 +1,312 @@ +// Copyright 2023 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "grpc/kalman_service.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "grpc/kalman.pb.h" +#include "mjpc/estimators/kalman.h" + +namespace kalman_grpc { + +using ::kalman::CovarianceRequest; +using ::kalman::CovarianceResponse; +using ::kalman::InitRequest; +using ::kalman::InitResponse; +using ::kalman::NoiseRequest; +using ::kalman::NoiseResponse; +using ::kalman::ResetRequest; +using ::kalman::ResetResponse; +using ::kalman::SettingsRequest; +using ::kalman::SettingsResponse; +using ::kalman::StateRequest; +using ::kalman::StateResponse; +using ::kalman::TimersRequest; +using ::kalman::TimersResponse; +using ::kalman::UpdateMeasurementRequest; +using ::kalman::UpdateMeasurementResponse; +using ::kalman::UpdatePredictionRequest; +using ::kalman::UpdatePredictionResponse; + +// TODO(taylor): make CheckSize utility function +namespace { +absl::Status CheckSize(std::string_view name, int model_size, int vector_size) { + std::ostringstream error_string; + if (model_size != vector_size) { + error_string << "expected " << name << " size " << model_size << ", got " + << vector_size; + return absl::InvalidArgumentError(error_string.str()); + } + return absl::OkStatus(); +} +} // namespace + +#define CHECK_SIZE(name, n1, n2) \ + { \ + auto expr = (CheckSize(name, n1, n2)); \ + if (!(expr).ok()) { \ + return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, \ + (expr).ToString()); \ + } \ + } + +KalmanService::~KalmanService() {} + +grpc::Status KalmanService::Init(grpc::ServerContext* context, + const kalman::InitRequest* request, + kalman::InitResponse* response) { + + // ----- initialize with model ----- // + mjpc::UniqueMjModel tmp_model = {nullptr, mj_deleteModel}; + + // convert message + if (request->has_model() && request->model().has_mjb()) { + std::string_view mjb = request->model().mjb(); + static constexpr char file[] = "temporary-filename.mjb"; + // mjVFS structs need to be allocated on the heap, because it's ~2MB + auto vfs = std::make_unique(); + mj_defaultVFS(vfs.get()); + mj_makeEmptyFileVFS(vfs.get(), file, mjb.size()); + int file_idx = mj_findFileVFS(vfs.get(), file); + memcpy(vfs->filedata[file_idx], mjb.data(), mjb.size()); + tmp_model = {mj_loadModel(file, vfs.get()), mj_deleteModel}; + mj_deleteFileVFS(vfs.get(), file); + } else if (request->has_model() && request->model().has_xml()) { + std::string_view model_xml = request->model().xml(); + char load_error[1024] = ""; + + // TODO(taylor): utilize grpc_agent_util method + static constexpr char file[] = "temporary-filename.xml"; + // mjVFS structs need to be allocated on the heap, because it's ~2MB + auto vfs = std::make_unique(); + mj_defaultVFS(vfs.get()); + mj_makeEmptyFileVFS(vfs.get(), file, model_xml.size()); + int file_idx = mj_findFileVFS(vfs.get(), file); + memcpy(vfs->filedata[file_idx], model_xml.data(), model_xml.size()); + tmp_model = {mj_loadXML(file, vfs.get(), load_error, sizeof(load_error)), + mj_deleteModel}; + mj_deleteFileVFS(vfs.get(), file); + } else { + mju_error("Failed to create mjModel."); + } + + // move + kalman_model_override_ = std::move(tmp_model); + mjModel* model = kalman_model_override_.get(); + + // initialize kalman + kalman_.Initialize(model); + kalman_.Reset(); + + return grpc::Status::OK; +} + +grpc::Status KalmanService::Reset(grpc::ServerContext* context, + const kalman::ResetRequest* request, + kalman::ResetResponse* response) { + if (!Initialized()) { + return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; + } + + // reset + kalman_.Reset(); + + return grpc::Status::OK; +} + +grpc::Status KalmanService::Settings(grpc::ServerContext* context, + const kalman::SettingsRequest* request, + kalman::SettingsResponse* response) { + if (!Initialized()) { + return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; + } + + // settings + kalman::Settings input = request->settings(); + kalman::Settings* output = response->mutable_settings(); + + // epsilon + if (input.has_epsilon()) { + kalman_.settings.epsilon = input.epsilon(); + } + output->set_epsilon(kalman_.settings.epsilon); + + // flg_centered + if (input.has_flg_centered()) { + kalman_.settings.flg_centered = input.flg_centered(); + } + output->set_flg_centered(kalman_.settings.flg_centered); + + return grpc::Status::OK; +} + +grpc::Status KalmanService::UpdateMeasurement( + grpc::ServerContext* context, const kalman::UpdateMeasurementRequest* request, + kalman::UpdateMeasurementResponse* response) { + if (!Initialized()) { + return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; + } + + // measurement update + kalman_.UpdateMeasurement(request->ctrl().data(), request->sensor().data()); + + return grpc::Status::OK; +} + +grpc::Status KalmanService::UpdatePrediction( + grpc::ServerContext* context, const kalman::UpdatePredictionRequest* request, + kalman::UpdatePredictionResponse* response) { + if (!Initialized()) { + return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; + } + + // prediction update + kalman_.UpdatePrediction(); + + return grpc::Status::OK; +} + +grpc::Status KalmanService::Timers(grpc::ServerContext* context, + const kalman::TimersRequest* request, + kalman::TimersResponse* response) { + if (!Initialized()) { + return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; + } + + // measurement + response->set_measurement(kalman_.TimerMeasurement()); + + // prediction + response->set_prediction(kalman_.TimerPrediction()); + + return grpc::Status::OK; +} + +grpc::Status KalmanService::State(grpc::ServerContext* context, + const kalman::StateRequest* request, + kalman::StateResponse* response) { + if (!Initialized()) { + return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; + } + + // unpack input/output + kalman::State input = request->state(); + kalman::State* output = response->mutable_state(); + + // set state + int nstate = kalman_.model->nq + kalman_.model->nv; + if (input.state_size() > 0) { + CHECK_SIZE("state", nstate, input.state_size()); + mju_copy(kalman_.state.data(), input.state().data(), nstate); + } + + // get state + double* state = kalman_.state.data(); + for (int i = 0; i < nstate; i++) { + output->add_state(state[i]); + } + + return grpc::Status::OK; +} + +grpc::Status KalmanService::Covariance(grpc::ServerContext* context, + const kalman::CovarianceRequest* request, + kalman::CovarianceResponse* response) { + if (!Initialized()) { + return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; + } + + // unpack input/output + kalman::Covariance input = request->covariance(); + kalman::Covariance* output = response->mutable_covariance(); + + // dimensions + int nvelocity = 2 * kalman_.model->nv; + int ncovariance = nvelocity * nvelocity; + + // set dimension + output->set_dimension(nvelocity); + + // set covariance + if (input.covariance_size() > 0) { + CHECK_SIZE("covariance", ncovariance, input.covariance_size()); + mju_copy(kalman_.covariance.data(), input.covariance().data(), ncovariance); + } + + // get covariance + double* covariance = kalman_.covariance.data(); + for (int i = 0; i < ncovariance; i++) { + output->add_covariance(covariance[i]); + } + + return grpc::Status::OK; +} + +grpc::Status KalmanService::Noise(grpc::ServerContext* context, + const kalman::NoiseRequest* request, + kalman::NoiseResponse* response) { + if (!Initialized()) { + return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; + } + + // unpack input/output + kalman::Noise input = request->noise(); + kalman::Noise* output = response->mutable_noise(); + + // dimensions + int nprocess = 2 * kalman_.model->nv; + int nsensor = kalman_.model->nsensordata; + + // set process noise + if (input.process_size() > 0) { + CHECK_SIZE("process noise", nprocess, input.process_size()); + mju_copy(kalman_.noise_process.data(), input.process().data(), nprocess); + } + + // get process noise + double* process = kalman_.noise_process.data(); + for (int i = 0; i < nprocess; i++) { + output->add_process(process[i]); + } + + // set sensor noise + if (input.sensor_size() > 0) { + CHECK_SIZE("sensor noise", nsensor, input.sensor_size()); + mju_copy(kalman_.noise_sensor.data(), input.sensor().data(), nsensor); + } + + // get sensor noise + double* sensor = kalman_.noise_sensor.data(); + for (int i = 0; i < nsensor; i++) { + output->add_sensor(sensor[i]); + } + + return grpc::Status::OK; +} + +#undef CHECK_SIZE + +} // namespace kalman_grpc \ No newline at end of file diff --git a/grpc/kalman_service.h b/grpc/kalman_service.h new file mode 100644 index 000000000..6853e4d71 --- /dev/null +++ b/grpc/kalman_service.h @@ -0,0 +1,90 @@ +// Copyright 2023 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// An implementation of the `Kalman` gRPC service. + +#ifndef GRPC_KALMAN_SERVICE_H +#define GRPC_KALMAN_SERVICE_H + +#include +#include +#include + +#include +#include + +#include "grpc/kalman.grpc.pb.h" +#include "grpc/kalman.pb.h" +#include "mjpc/estimators/kalman.h" +#include "mjpc/threadpool.h" +#include "mjpc/utilities.h" + +namespace kalman_grpc { + +class KalmanService final : public kalman::Kalman::Service { + public: + explicit KalmanService() : thread_pool_(1) {} + ~KalmanService(); + + grpc::Status Init(grpc::ServerContext* context, + const kalman::InitRequest* request, + kalman::InitResponse* response) override; + + grpc::Status Reset(grpc::ServerContext* context, + const kalman::ResetRequest* request, + kalman::ResetResponse* response) override; + + grpc::Status Settings(grpc::ServerContext* context, + const kalman::SettingsRequest* request, + kalman::SettingsResponse* response) override; + + grpc::Status UpdateMeasurement( + grpc::ServerContext* context, + const kalman::UpdateMeasurementRequest* request, + kalman::UpdateMeasurementResponse* response) override; + + grpc::Status UpdatePrediction( + grpc::ServerContext* context, const kalman::UpdatePredictionRequest* request, + kalman::UpdatePredictionResponse* response) override; + + grpc::Status Timers(grpc::ServerContext* context, + const kalman::TimersRequest* request, + kalman::TimersResponse* response) override; + + grpc::Status State(grpc::ServerContext* context, + const kalman::StateRequest* request, + kalman::StateResponse* response) override; + + grpc::Status Covariance(grpc::ServerContext* context, + const kalman::CovarianceRequest* request, + kalman::CovarianceResponse* response) override; + + grpc::Status Noise(grpc::ServerContext* context, + const kalman::NoiseRequest* request, + kalman::NoiseResponse* response) override; + + private: + bool Initialized() const { return kalman_.model; } + + // kalman + mjpc::Kalman kalman_; + mjpc::UniqueMjModel kalman_model_override_ = {nullptr, mj_deleteModel}; + + // threadpool + mjpc::ThreadPool thread_pool_; +}; + +} // namespace kalman_grpc + +#endif // GRPC_KALMAN_SERVICE_H \ No newline at end of file diff --git a/mjpc/CMakeLists.txt b/mjpc/CMakeLists.txt index fa336dbe9..ede20375f 100644 --- a/mjpc/CMakeLists.txt +++ b/mjpc/CMakeLists.txt @@ -100,9 +100,13 @@ add_library( planners/ilqs/planner.h estimators/batch.cc estimators/batch.h + estimators/estimator.h + estimators/gui.cc + estimators/gui.h estimators/include.cc estimators/include.h - estimators/estimator.h + estimators/kalman.cc + estimators/kalman.h estimators/trajectory.h app.cc app.h diff --git a/mjpc/estimators/batch.cc b/mjpc/estimators/batch.cc index 0476ef7ea..fd7e133db 100644 --- a/mjpc/estimators/batch.cc +++ b/mjpc/estimators/batch.cc @@ -17,6 +17,8 @@ #include #include "mjpc/array_safety.h" +#include "mjpc/estimators/estimator.h" +#include "mjpc/estimators/gui.h" #include "mjpc/norm.h" #include "mjpc/threadpool.h" #include "mjpc/utilities.h" @@ -80,6 +82,11 @@ void Batch::Initialize(const mjModel* model) { configuration_length_ = GetNumberOrDefault(3, model, "batch_configuration_length"); + // check configuration length + if (configuration_length_ > max_history_) { + mju_error("configuration_length > max_history: increase max history\n"); + } + // -- trajectories -- // configuration.Initialize(nq, configuration_length_); velocity.Initialize(nv, configuration_length_); @@ -147,7 +154,8 @@ void Batch::Initialize(const mjModel* model) { block_force_velocity_.Initialize(nv * nv, configuration_length_); block_force_acceleration_.Initialize(nv * nv, configuration_length_); - block_force_previous_configuration_.Initialize(nv * nv, configuration_length_); + block_force_previous_configuration_.Initialize(nv * nv, + configuration_length_); block_force_current_configuration_.Initialize(nv * nv, configuration_length_); block_force_next_configuration_.Initialize(nv * nv, configuration_length_); block_force_configurations_.Initialize(nv * 3 * nv, configuration_length_); @@ -251,6 +259,14 @@ void Batch::Initialize(const mjModel* model) { scratch0_covariance_.resize((nv * max_history_) * (nv * max_history_)); scratch1_covariance_.resize((nv * max_history_) * (nv * max_history_)); + // conditioned matrix + mat00_.resize((nv * max_history_) * (nv * max_history_)); + mat10_.resize((nv * max_history_) * (nv * max_history_)); + mat11_.resize((nv * max_history_) * (nv * max_history_)); + condmat_.resize((nv * max_history_) * (nv * max_history_)); + scratch0_condmat_.resize((nv * max_history_) * (nv * max_history_)); + scratch1_condmat_.resize((nv * max_history_) * (nv * max_history_)); + // regularization regularization_ = settings.regularization_initial; @@ -272,19 +288,46 @@ void Batch::Initialize(const mjModel* model) { // settings settings.band_prior = (bool)GetNumberOrDefault(1, model, "batch_band_covariance"); + + // -- trajectory cache -- // + configuration_cache_.Initialize(nq, configuration_length_); + velocity_cache_.Initialize(nv, configuration_length_); + acceleration_cache_.Initialize(nv, configuration_length_); + act_cache_.Initialize(na, configuration_length_); + times_cache_.Initialize(1, configuration_length_); + + // ctrl + ctrl_cache_.Initialize(model->nu, configuration_length_); + + // prior + configuration_previous_cache_.Initialize(nq, configuration_length_); + + // sensor + sensor_measurement_cache_.Initialize(nsensordata_, configuration_length_); + sensor_prediction_cache_.Initialize(nsensordata_, configuration_length_); + sensor_mask_cache_.Initialize(nsensor_, configuration_length_); + + // force + force_measurement_cache_.Initialize(nv, configuration_length_); + force_prediction_cache_.Initialize(nv, configuration_length_); } // reset memory -void Batch::Reset() { +void Batch::Reset(const mjData* data) { // dimension int nq = model->nq, nv = model->nv, na = model->na; // data mjData* d = data_[0].get(); - // set home keyframe - int home_id = mj_name2id(model, mjOBJ_KEY, "home"); - if (home_id >= 0) mj_resetDataKeyframe(model, d, home_id); + if (data) { + // copy input data + mj_copyData(d, model, data); + } else { + // set home keyframe + int home_id = mj_name2id(model, mjOBJ_KEY, "home"); + if (home_id >= 0) mj_resetDataKeyframe(model, d, home_id); + } // forward evaluation mj_forward(model, d); @@ -293,8 +336,7 @@ void Batch::Reset() { mju_copy(state.data(), d->qpos, nq); mju_copy(state.data() + nq, d->qvel, nv); mju_copy(state.data() + nq + nv, d->act, na); - d->time = 0.0; - time = 0.0; + time = d->time; // covariance mju_eye(covariance.data(), ndstate_); @@ -313,6 +355,9 @@ void Batch::Reset() { GetNumberOrDefault(1.0e-4, model, "estimator_sensor_noise_scale"); std::fill(noise_sensor.begin(), noise_sensor.end(), noise_sensor_scl); + // scale prior + scale_prior = GetNumberOrDefault(1.0e-1, model, "batch_scale_prior"); + // trajectories configuration.Reset(); velocity.Reset(); @@ -330,8 +375,8 @@ void Batch::Reset() { // sensor mask sensor_mask.Reset(); - for (int i = 0; i < nsensor_ * configuration_length_; i++) { - sensor_mask.Data()[i] = 1; // sensor on + for (int i = 0; i < nsensor_ * max_history_; i++) { + sensor_mask.Data()[i] = 1; // sensor on } // force @@ -455,6 +500,14 @@ void Batch::Reset() { std::fill(scratch0_covariance_.begin(), scratch0_covariance_.end(), 0.0); std::fill(scratch1_covariance_.begin(), scratch1_covariance_.end(), 0.0); + // conditioned matrix + std::fill(mat00_.begin(), mat00_.end(), 0.0); + std::fill(mat10_.begin(), mat10_.end(), 0.0); + std::fill(mat11_.begin(), mat11_.end(), 0.0); + std::fill(condmat_.begin(), condmat_.end(), 0.0); + std::fill(scratch0_condmat_.begin(), scratch0_condmat_.end(), 0.0); + std::fill(scratch1_condmat_.begin(), scratch1_condmat_.end(), 0.0); + // timer std::fill(timer_.prior_step.begin(), timer_.prior_step.end(), 0.0); std::fill(timer_.sensor_step.begin(), timer_.sensor_step.end(), 0.0); @@ -469,79 +522,33 @@ void Batch::Reset() { cost_count_ = 0; solve_status_ = kUnsolved; - // -- initialize -- // - if (settings.filter) { - settings.gradient_tolerance = 1.0e-8; - settings.max_smoother_iterations = 1; - settings.max_search_iterations = 10; - - // timestep - double timestep = model->opt.timestep; - - // set q1 - configuration.Set(state.data(), 1); - configuration_previous.Set(configuration.Get(1), 1); - - // set q0 - double* q0 = configuration.Get(0); - mju_copy(q0, state.data(), nq); - mj_integratePos(model, q0, state.data() + nq, -1.0 * timestep); - configuration_previous.Set(configuration.Get(0), 0); - - // set times - double current_time = -1.0 * timestep; - times.Set(¤t_time, 0); - for (int i = 1; i < configuration_length_; i++) { - // increment time - current_time += timestep; - - // set - times.Set(¤t_time, i); - } - - // -- set initial position-based measurements -- // - - // data - mjData* data = data_[0].get(); + // initialize filter + if (settings.filter) InitializeFilter(); - // set q0 - mju_copy(data->qpos, q0, model->nq); + // trajectory cache + configuration_cache_.Reset(); + velocity_cache_.Reset(); + acceleration_cache_.Reset(); + act_cache_.Reset(); + times_cache_.Reset(); + ctrl_cache_.Reset(); - // evaluate position - mj_fwdPosition(model, data); - mj_sensorPos(model, data); - - // y0 - double* y0 = sensor_measurement.Get(0); - mju_zero(y0, nsensordata_); - - // loop over sensors - for (int i = 0; i < nsensor_; i++) { - // measurement sensor index - int index = sensor_start_ + i; - - // need stage - int sensor_stage = model->sensor_needstage[index]; - - if (sensor_stage == mjSTAGE_POS) { - // address - int sensor_adr = model->sensor_adr[index]; - - // dimension - int sensor_dim = model->sensor_dim[index]; + // prior + configuration_previous_cache_.Reset(); - // set sensor - mju_copy(y0 + sensor_adr - sensor_start_index_, - data->sensordata + sensor_adr, sensor_dim); - } - } + // sensor + sensor_measurement_cache_.Reset(); + sensor_prediction_cache_.Reset(); - // prior weight (skip act) - for (int i = 0; i < ndstate_ - na; i++) { - weight_prior[nv * configuration_length_ * i + i] = - 1.0 / covariance[ndstate_ * i + i]; - } + // sensor mask + sensor_mask_cache_.Reset(); + for (int i = 0; i < nsensor_ * configuration_length_; i++) { + sensor_mask_cache_.Data()[i] = 1; // sensor on } + + // force + force_measurement_cache_.Reset(); + force_prediction_cache_.Reset(); } // update @@ -549,32 +556,25 @@ void Batch::Update(const double* ctrl, const double* sensor) { // start timer auto start = std::chrono::steady_clock::now(); - // check configuration length - if (configuration_length_ != 3) { - mju_error("batch filter only supports configuration length = 3\n"); - } - // dimensions int nq = model->nq, nv = model->nv, na = model->na, nu = model->nu; + // data + mjData* d = data_[0].get(); + // current time index - int t = 1; + int t = current_time_index; // configurations double* q0 = configuration.Get(t - 1); double* q1 = configuration.Get(t); - // data - mjData* d = data_[0].get(); - // -- next qpos -- // // set state mju_copy(d->qpos, q1, model->nq); mj_differentiatePos(model, d->qvel, model->opt.timestep, q0, q1); - // TODO(taylor): set time - // set ctrl mju_copy(d->ctrl, ctrl, nu); @@ -582,6 +582,7 @@ void Batch::Update(const double* ctrl, const double* sensor) { mj_step(model, d); // -- set batch data -- // + // set next qpos configuration.Set(d->qpos, t + 1); configuration_previous.Set(d->qpos, t + 1); @@ -598,7 +599,17 @@ void Batch::Update(const double* ctrl, const double* sensor) { // set force measurement force_measurement.Set(d->qfrc_actuator, t); - // measurement update + // -- measurement update -- // + + // cache configuration length + int configuration_length_cache = configuration_length_; + + // set reduced configuration length for optimization + configuration_length_ = current_time_index + 2; + if (configuration_length_ != configuration_length_cache) { + ShiftResizeTrajectory(0, configuration_length_); + } + // TODO(taylor): preallocate pool ThreadPool pool(1); Optimize(pool); @@ -607,13 +618,77 @@ void Batch::Update(const double* ctrl, const double* sensor) { mju_copy(state.data(), configuration.Get(t + 1), nq); mju_copy(state.data() + nq, velocity.Get(t + 1), nv); mju_copy(state.data() + nq + nv, act.Get(t + 1), na); - time = d->time; + time = times.Get(t + 1)[0]; + + // -- update prior weights -- // + + // dimension + int nvar = nv * configuration_length_; + + // prior weights + double* weights = weight_prior.data(); + + // recursive update + if (settings.recursive_prior_update && + configuration_length_ == configuration_length_cache) { + // condition dimension + int ncondition = nv * (configuration_length_ - 1); + + // compute conditioned matrix + double* condmat = condmat_.data(); + ConditionMatrix(condmat, cost_hessian.data(), mat00_.data(), mat10_.data(), + mat11_.data(), scratch0_condmat_.data(), + scratch1_condmat_.data(), nvar, nv, ncondition); + + // zero memory + mju_zero(weights, nvar * nvar); + + // set conditioned matrix in prior weights + SetBlockInMatrix(weights, condmat, 1.0, nvar, nvar, ncondition, ncondition, + 0, 0); + + // set bottom right to scale_prior * I + for (int i = ncondition; i < nvar; i++) { + weights[nvar * i + i] = scale_prior; + } + } else { + // dimension + int nvar_new = nvar; + if (current_time_index < configuration_length_ - 2) { + nvar_new += nv; + } - // update prior weights - // TODO(taylor) + // check same size + if (nvar != nvar_new) { + // get previous weights + double* previous_weights = scratch0_condmat_.data(); + mju_copy(previous_weights, weights, nvar * nvar); - // shift trajectories - Shift(1); + // set previous in new weights (dimension may have increased) + mju_zero(weights, nvar_new * nvar_new); + SetBlockInMatrix(weights, previous_weights, 1.0, nvar_new, nvar_new, nvar, + nvar, 0, 0); + + // scale_prior * I + for (int i = nvar; i < nvar_new; i++) { + weights[nvar_new * i + i] = scale_prior; + } + } + } + + // restore configuration length + if (configuration_length_ != configuration_length_cache) { + ShiftResizeTrajectory(0, configuration_length_cache); + } + configuration_length_ = configuration_length_cache; + + // check estimation horizon + if (current_time_index < configuration_length_ - 2) { + current_time_index++; + } else { + // shift trajectories once estimation horizon is filled + Shift(1); + } // stop timer timer_.update = 1.0e-3 * GetDuration(start); @@ -1026,7 +1101,7 @@ double Batch::CostSensor(double* gradient, double* hessian) { if (gradient) mju_zero(gradient, nvar); if (hessian) mju_zero(hessian, nvar * nvar); - // time scaling + // time scaling double time_scale = 1.0; double time_scale2 = 1.0; if (settings.time_scaling_sensor) { @@ -1047,15 +1122,15 @@ double Batch::CostSensor(double* gradient, double* hessian) { // unpack block double* block; - int block_columns; - if (t == 0) { // only position sensors + int block_columns; + if (t == 0) { // only position sensors block = block_sensor_configuration_.Get(t) + sensor_start_index_ * nv; block_columns = nv; - } else { // position, velocity, acceleration sensors + } else { // position, velocity, acceleration sensors block = block_sensor_configurations_.Get(t); block_columns = 3 * nv; } - + // shift int shift_sensor = 0; @@ -1064,10 +1139,10 @@ double Batch::CostSensor(double* gradient, double* hessian) { // start cost timer auto start_cost = std::chrono::steady_clock::now(); - // sensor stage + // sensor stage int sensor_stage = model->sensor_needstage[sensor_start_ + i]; - - // time scaling weight + + // time scaling weight double time_weight = 1.0; if (sensor_stage == mjSTAGE_VEL) { time_weight = time_scale; @@ -1085,7 +1160,8 @@ double Batch::CostSensor(double* gradient, double* hessian) { double* rti = rt + shift_sensor; // weight - double weight = time_weight / noise_sensor[i] / nsi / (configuration_length_ - 1); + double weight = + time_weight / noise_sensor[i] / nsi / (configuration_length_ - 1); // parameters double* pi = norm_parameters_sensor.data() + kMaxNormParameters * i; @@ -1135,8 +1211,8 @@ double Batch::CostSensor(double* gradient, double* hessian) { block_columns); // add - mju_addToScl(gradient + nv * std::max(0, t - 1), scratch0_sensor_.data(), weight, - block_columns); + mju_addToScl(gradient + nv * std::max(0, t - 1), + scratch0_sensor_.data(), weight, block_columns); } // Hessian (Gauss-Newton): drdq' * d2ndr2 * drdq @@ -1153,8 +1229,9 @@ double Batch::CostSensor(double* gradient, double* hessian) { mju_mulMatTMat(tmp1, blocki, tmp0, nsi, block_columns, block_columns); // add - AddBlockInMatrix(hessian, tmp1, weight, nvar, nvar, block_columns, block_columns, - nv * std::max(0, t - 1), nv * std::max(0, t - 1)); + AddBlockInMatrix(hessian, tmp1, weight, nvar, nvar, block_columns, + block_columns, nv * std::max(0, t - 1), + nv * std::max(0, t - 1)); } // shift by individual sensor dimension @@ -1191,7 +1268,7 @@ double Batch::CostForce(double* gradient, double* hessian) { if (gradient) mju_zero(gradient, nvar); if (hessian) mju_zero(hessian, nvar * nvar); - // time scaling + // time scaling double time_scale2 = 1.0; if (settings.time_scaling_force) { time_scale2 = model->opt.timestep * model->opt.timestep; @@ -1220,7 +1297,8 @@ double Batch::CostForce(double* gradient, double* hessian) { // quadratic cost for (int i = 0; i < nv; i++) { // weight - double weight = time_scale2 / noise_process[i] / nv / (configuration_length_ - 2); + double weight = + time_scale2 / noise_process[i] / nv / (configuration_length_ - 2); // gradient norm_gradient[i] = weight * rt[i]; @@ -1254,7 +1332,8 @@ double Batch::CostForce(double* gradient, double* hessian) { mju_mulMatTVec(scratch0_force_.data(), block, norm_gradient, nv, 3 * nv); // add - mju_addToScl(gradient + (t - 1) * nv, scratch0_force_.data(), 1.0, 3 * nv); + mju_addToScl(gradient + (t - 1) * nv, scratch0_force_.data(), 1.0, + 3 * nv); } // Hessian (Gauss-Newton): drdq' * d2ndr2 * drdq @@ -1268,8 +1347,8 @@ double Batch::CostForce(double* gradient, double* hessian) { mju_mulMatTMat(tmp1, block, tmp0, nv, 3 * nv, 3 * nv); // add - AddBlockInMatrix(hessian, tmp1, 1.0, nvar, nvar, 3 * nv, 3 * nv, nv * (t - 1), - nv * (t - 1)); + AddBlockInMatrix(hessian, tmp1, 1.0, nvar, nvar, 3 * nv, 3 * nv, + nv * (t - 1), nv * (t - 1)); } } @@ -1294,7 +1373,7 @@ void Batch::ResidualSensor() { // sensor difference mju_sub(rt, yt_model, yt_sensor, nsensordata_); - // zero out non-position sensors at first time step + // zero out non-position sensors at first time step if (t == 0) { // loop over position sensors for (int i = 0; i < nsensor_; i++) { @@ -1417,7 +1496,8 @@ void Batch::BlockSensor(int index) { if (settings.assemble_sensor_jacobian) { // set block SetBlockInMatrix(jacobian_sensor_.data(), dsdq012, 1.0, nsen, nvar, - nsensordata_, 3 * nv, index * nsensordata_, (index - 1) * nv); + nsensordata_, 3 * nv, index * nsensordata_, + (index - 1) * nv); } } @@ -1580,18 +1660,18 @@ void Batch::InverseDynamicsPrediction(ThreadPool& pool) { // first time step pool.Schedule([&batch = *this, nq, nv, nu]() { - // time index + // time index int t = 0; // data mjData* d = batch.data_[t].get(); - // terms + // terms double* q0 = batch.configuration.Get(t); double* y0 = batch.sensor_prediction.Get(t); mju_zero(y0, batch.nsensordata_); - // set data + // set data mju_copy(d->qpos, q0, nq); mju_zero(d->qvel, nv); mju_zero(d->ctrl, nu); @@ -1601,17 +1681,17 @@ void Batch::InverseDynamicsPrediction(ThreadPool& pool) { mj_fwdPosition(batch.model, d); mj_sensorPos(batch.model, d); - // loop over position sensors + // loop over position sensors for (int i = 0; i < batch.nsensor_; i++) { // sensor stage int sensor_stage = batch.model->sensor_needstage[batch.sensor_start_ + i]; - // check for position + // check for position if (sensor_stage == mjSTAGE_POS) { - // dimension + // dimension int sensor_dim = batch.model->sensor_dim[batch.sensor_start_ + i]; - // address + // address int sensor_adr = batch.model->sensor_adr[batch.sensor_start_ + i]; // copy sensor data @@ -1678,17 +1758,17 @@ void Batch::InverseDynamicsDerivatives(ThreadPool& pool) { // first time step pool.Schedule([&batch = *this, nq, nv, nu]() { - // time index + // time index int t = 0; // data mjData* d = batch.data_[t].get(); - // terms + // terms double* q0 = batch.configuration.Get(t); double* dsdq = batch.block_sensor_configuration_.Get(t); - // set data + // set data mju_copy(d->qpos, q0, nq); mju_zero(d->qvel, nv); mju_zero(d->ctrl, nu); @@ -1884,6 +1964,159 @@ void Batch::VelocityAccelerationDerivatives() { timer_.velacc_derivatives += GetDuration(start); } +// initialize filter mode +void Batch::InitializeFilter() { + // dimensions + int nq = model->nq, nv = model->nv; + + // filter mode status + current_time_index = 1; + + // filter settings + settings.gradient_tolerance = 1.0e-6; + settings.max_smoother_iterations = 1; + settings.max_search_iterations = 10; + settings.recursive_prior_update = true; + + // timestep + double timestep = model->opt.timestep; + + // set q1 + configuration.Set(state.data(), 1); + configuration_previous.Set(configuration.Get(1), 1); + + // set q0 + double* q0 = configuration.Get(0); + mju_copy(q0, state.data(), nq); + mj_integratePos(model, q0, state.data() + nq, -1.0 * timestep); + configuration_previous.Set(configuration.Get(0), 0); + + // set times + double current_time = time - timestep; + times.Set(¤t_time, 0); + for (int i = 1; i < configuration_length_; i++) { + // increment time + current_time += timestep; + + // set + times.Set(¤t_time, i); + } + + // -- set initial position-based measurements -- // + + // data + mjData* data = data_[0].get(); + + // set q0 + mju_copy(data->qpos, q0, model->nq); + + // evaluate position + mj_fwdPosition(model, data); + mj_sensorPos(model, data); + + // y0 + double* y0 = sensor_measurement.Get(0); + mju_zero(y0, nsensordata_); + + // loop over sensors + for (int i = 0; i < nsensor_; i++) { + // measurement sensor index + int index = sensor_start_ + i; + + // need stage + int sensor_stage = model->sensor_needstage[index]; + + // position sensor + if (sensor_stage == mjSTAGE_POS) { + // address + int sensor_adr = model->sensor_adr[index]; + + // dimension + int sensor_dim = model->sensor_dim[index]; + + // set sensor + mju_copy(y0 + sensor_adr - sensor_start_index_, + data->sensordata + sensor_adr, sensor_dim); + } + } + + // prior weight + int nvar = nv * configuration_length_; + for (int i = 0; i < nvar; i++) { + weight_prior[nvar * i + i] = scale_prior; + } +} + +// shift head and resize trajectories +void Batch::ShiftResizeTrajectory(int new_head, int new_length) { + // reset cache + configuration_cache_.Reset(); + configuration_previous_cache_.Reset(); + velocity_cache_.Reset(); + acceleration_cache_.Reset(); + act_cache_.Reset(); + times_cache_.Reset(); + ctrl_cache_.Reset(); + sensor_measurement_cache_.Reset(); + sensor_prediction_cache_.Reset(); + sensor_mask_cache_.Reset(); + force_measurement_cache_.Reset(); + force_prediction_cache_.Reset(); + + // -- set cache length -- // + int length = configuration_length_; + + configuration_cache_.SetLength(length); + configuration_previous_cache_.SetLength(length); + velocity_cache_.SetLength(length); + acceleration_cache_.SetLength(length); + act_cache_.SetLength(length); + times_cache_.SetLength(length); + ctrl_cache_.SetLength(length); + sensor_measurement_cache_.SetLength(length); + sensor_prediction_cache_.SetLength(length); + sensor_mask_cache_.SetLength(length); + force_measurement_cache_.SetLength(length); + force_prediction_cache_.SetLength(length); + + // copy data to cache + for (int i = 0; i < length; i++) { + configuration_cache_.Set(configuration.Get(i), i); + configuration_previous_cache_.Set(configuration_previous.Get(i), i); + velocity_cache_.Set(velocity.Get(i), i); + acceleration_cache_.Set(acceleration.Get(i), i); + act_cache_.Set(act.Get(i), i); + times_cache_.Set(times.Get(i), i); + ctrl_cache_.Set(ctrl.Get(i), i); + sensor_measurement_cache_.Set(sensor_measurement.Get(i), i); + sensor_prediction_cache_.Set(sensor_prediction.Get(i), i); + sensor_mask_cache_.Set(sensor_mask.Get(i), i); + force_measurement_cache_.Set(force_measurement.Get(i), i); + force_prediction_cache_.Set(force_prediction.Get(i), i); + } + + // set configuration length + SetConfigurationLength(new_length); + + // set trajectory data + int length_copy = std::min(length, new_length); + for (int i = 0; i < length_copy; i++) { + configuration.Set(configuration_cache_.Get(new_head + i), i); + configuration_previous.Set(configuration_previous_cache_.Get(new_head + i), + i); + velocity.Set(velocity_cache_.Get(new_head + i), i); + acceleration.Set(acceleration_cache_.Get(new_head + i), i); + act.Set(act_cache_.Get(new_head + i), i); + times.Set(times_cache_.Get(new_head + i), i); + ctrl.Set(ctrl_cache_.Get(new_head + i), i); + sensor_measurement.Set(sensor_measurement_cache_.Get(new_head + i), i); + sensor_prediction.Set(sensor_prediction_cache_.Get(new_head + i), i); + sensor_mask.Set(sensor_mask_cache_.Get(new_head + i), i); + force_measurement.Set(force_measurement_cache_.Get(new_head + i), i); + force_prediction.Set(force_prediction_cache_.Get(new_head + i), i); + } +} + // compute total cost // TODO(taylor): fix timers double Batch::Cost(double* gradient, double* hessian, ThreadPool& pool) { @@ -2100,7 +2333,7 @@ void Batch::Optimize(ThreadPool& pool) { // -- search direction -- // // check regularization - if (regularization_ >= MAX_REGULARIZATION - 1.0e-6) { + if (regularization_ >= kMaxBatchRegularization - 1.0e-6) { // set solve status solve_status_ = kMaxRegularizationFailure; @@ -2230,12 +2463,12 @@ void Batch::Optimize(ThreadPool& pool) { if (reduction_ratio_ > 0.75) { // decrease regularization_ = - mju_max(MIN_REGULARIZATION, + mju_max(kMinBatchRegularization, regularization_ / settings.regularization_scaling); } else if (reduction_ratio_ < 0.25) { // increase regularization_ = - mju_min(MAX_REGULARIZATION, + mju_min(kMaxBatchRegularization, regularization_ * settings.regularization_scaling); } } @@ -2291,13 +2524,14 @@ void Batch::SearchDirection() { double min_diag = 0.0; while (min_diag <= 0.0) { // failure - if (regularization_ >= MAX_REGULARIZATION) { + if (regularization_ >= kMaxBatchRegularization) { printf("min diag = %f\n", min_diag); mju_error("cost Hessian factorization failure: MAX REGULARIZATION\n"); } // copy - mju_copy(hessian_band_factor, hessian_band, ntotal * ntotal); // TODO(taylor): band copy + mju_copy(hessian_band_factor, hessian_band, + ntotal * ntotal); // TODO(taylor): band copy // factorize min_diag = mju_cholFactorBand(hessian_band_factor, ntotal, nband, ndense, @@ -2320,7 +2554,7 @@ void Batch::SearchDirection() { int rank = 0; while (rank < ntotal) { // failure - if (regularization_ >= MAX_REGULARIZATION) { + if (regularization_ >= kMaxBatchRegularization) { mju_error("cost Hessian factorization failure: MAX REGULARIZATION\n"); } @@ -2348,6 +2582,13 @@ void Batch::SearchDirection() { // search direction norm search_direction_norm_ = InfinityNorm(direction, ntotal); + // set regularization + if (regularization_ > 0.0) { + for (int i = 0; i < ntotal; i++) { + hessian[i * ntotal + i] += regularization_; + } + } + // end timer timer_.search_direction += GetDuration(search_direction_start); } @@ -2518,7 +2759,7 @@ std::string StatusString(int code) { case kSmallDirectionFailure: return "SMALL_DIRECTION_FAILURE"; case kMaxRegularizationFailure: - return "MAX_REGULARIZATION_FAILURE"; + return "kMaxBatchRegularization_FAILURE"; case kCostDifferenceFailure: return "COST_DIFFERENCE_FAILURE"; case kExpectedDecreaseFailure: @@ -2531,18 +2772,22 @@ std::string StatusString(int code) { } // estimator-specific GUI elements -void Batch::GUI(mjUI& ui, double* process_noise, double* sensor_noise, - double& timestep, int& integrator) { +void Batch::GUI(mjUI& ui, EstimatorGUIData& data) { // ----- estimator ------ // mjuiDef defEstimator[] = { {mjITEM_SECTION, "Estimator Settings", 1, nullptr, "AP"}, // needs new section to satisfy mjMAXUIITEM {mjITEM_BUTTON, "Reset", 2, nullptr, ""}, - {mjITEM_SLIDERNUM, "Timestep", 2, ×tep, "1.0e-3 0.1"}, - {mjITEM_SELECT, "Integrator", 2, &integrator, + {mjITEM_SLIDERNUM, "Timestep", 2, &data.timestep, "1.0e-3 0.1"}, + {mjITEM_SELECT, "Integrator", 2, &data.integrator, "Euler\nRK4\nImplicit\nFastImplicit"}, + {mjITEM_SLIDERINT, "Horizon", 2, &data.horizon, "3 3"}, {mjITEM_END}}; + // set estimation horizon limits + mju::sprintf_arr(defEstimator[4].other, "%i %i", kMinBatchHistory, + kMaxFilterHistory); + // add estimator mjui_add(&ui, defEstimator); @@ -2558,8 +2803,8 @@ void Batch::GUI(mjUI& ui, double* process_noise, double* sensor_noise, // add UI elements for (int i = 0; i < nv; i++) { // element - defProcessNoise[process_noise_shift] = {mjITEM_SLIDERNUM, "", 2, - process_noise + i, "1.0e-8 0.01"}; + defProcessNoise[process_noise_shift] = { + mjITEM_SLIDERNUM, "", 2, data.process_noise.data() + i, "1.0e-8 0.01"}; // set name mju::strcpy_arr(defProcessNoise[process_noise_shift].name, ""); @@ -2677,9 +2922,9 @@ void Batch::GUI(mjUI& ui, double* process_noise, double* sensor_noise, model->name_sensoradr[sensor_start_ + i]); // element - defSensorNoise[sensor_noise_shift] = {mjITEM_SLIDERNUM, "", 2, - sensor_noise + sensor_noise_shift - 1, - "1.0e-8 0.01"}; + defSensorNoise[sensor_noise_shift] = { + mjITEM_SLIDERNUM, "", 2, + data.sensor_noise.data() + sensor_noise_shift - 1, "1.0e-8 0.01"}; // sensor name sensor_str = name_sensor; @@ -2699,6 +2944,69 @@ void Batch::GUI(mjUI& ui, double* process_noise, double* sensor_noise, mjui_add(&ui, defSensorNoise); } +// set GUI data +void Batch::SetGUIData(EstimatorGUIData& data) { + mju_copy(noise_process.data(), data.process_noise.data(), DimensionProcess()); + mju_copy(noise_sensor.data(), data.sensor_noise.data(), DimensionSensor()); + model->opt.timestep = data.timestep; + model->opt.integrator = data.integrator; + + // store estimation horizon + int horizon = data.horizon; + + // changing horizon cases + if (horizon > configuration_length_) { // increase horizon + // -- prior weights resize -- // + int nvar = model->nv * configuration_length_; + int nvar_new = model->nv * horizon; + + // get previous weights + double* weights = weight_prior.data(); + double* previous_weights = scratch0_condmat_.data(); + mju_copy(previous_weights, weights, nvar * nvar); + + // set previous in new weights (dimension may have increased) + mju_zero(weights, nvar_new * nvar_new); + SetBlockInMatrix(weights, previous_weights, 1.0, nvar_new, nvar_new, nvar, + nvar, 0, 0); + + // scale_prior * I + for (int i = nvar; i < nvar_new; i++) { + weights[nvar_new * i + i] = scale_prior; + } + + // modify trajectories + ShiftResizeTrajectory(0, horizon); + + // update configuration length + configuration_length_ = horizon; + } else if (horizon < configuration_length_) { // decrease horizon + // -- prior weights resize -- // + int nvar = model->nv * configuration_length_; + int nvar_new = model->nv * horizon; + + // get previous weights + double* weights = weight_prior.data(); + double* previous_weights = scratch0_condmat_.data(); + BlockFromMatrix(previous_weights, weights, nvar_new, nvar_new, nvar, nvar, + 0, 0); + + // set previous in new weights (dimension may have increased) + mju_zero(weights, nvar * nvar); + mju_copy(weights, previous_weights, nvar_new * nvar_new); + + // compute difference in estimation horizons + int horizon_diff = configuration_length_ - horizon; + + // modify trajectories + ShiftResizeTrajectory(horizon_diff, horizon); + + // update configuration length and current time index + configuration_length_ = horizon; + current_time_index -= horizon_diff; + } +}; + // estimator-specific plots void Batch::Plots(mjvFigure* fig_planner, mjvFigure* fig_timer, int planner_shift, int timer_shift, int planning, @@ -2715,7 +3023,8 @@ void Batch::Plots(mjvFigure* fig_planner, mjvFigure* fig_timer, // mju_log10(trace), 100, planner_shift + 0, 0, 1, -100); // // legend - // mju::strcpy_arr(fig_planner->linename[planner_shift + 0], "Covariance Trace"); + // mju::strcpy_arr(fig_planner->linename[planner_shift + 0], "Covariance + // Trace"); // Batch timers double timer_bounds[2] = {0.0, 1.0}; @@ -2731,7 +3040,7 @@ void Batch::Plots(mjvFigure* fig_planner, mjvFigure* fig_timer, // increase regularization void Batch::IncreaseRegularization() { - regularization_ = mju_min(MAX_REGULARIZATION, + regularization_ = mju_min(kMaxBatchRegularization, regularization_ * settings.regularization_scaling); } diff --git a/mjpc/estimators/batch.h b/mjpc/estimators/batch.h index eb9dbdcb0..490bebf8e 100644 --- a/mjpc/estimators/batch.h +++ b/mjpc/estimators/batch.h @@ -21,6 +21,7 @@ #include #include "mjpc/estimators/estimator.h" +#include "mjpc/estimators/gui.h" #include "mjpc/estimators/trajectory.h" #include "mjpc/norm.h" #include "mjpc/threadpool.h" @@ -31,6 +32,8 @@ namespace mjpc { // defaults inline constexpr int kMinBatchHistory = 3; // minimum configuration trajectory length +inline constexpr int kMaxFilterHistory = + 32; // maximum batch filter estimation horizon // batch estimator status enum BatchStatus : int { @@ -52,8 +55,8 @@ enum SearchType : int { }; // maximum / minimum regularization -const double MAX_REGULARIZATION = 1.0e12; -const double MIN_REGULARIZATION = 1.0e-12; +inline constexpr double kMaxBatchRegularization = 1.0e12; +inline constexpr double kMinBatchRegularization = 1.0e-12; // ----- batch estimator ----- // // based on: "Physically-Consistent Sensor Fusion in Contact-Rich Behaviors" @@ -61,7 +64,10 @@ class Batch : public Estimator { public: // constructor Batch() = default; - Batch(int mode) { settings.filter = mode; } + Batch(int mode) { + settings.filter = mode; + max_history_ = kMaxFilterHistory; + } Batch(const mjModel* model, int length = 3, int max_history_ = 0) { // set max history length this->max_history_ = (max_history_ == 0 ? length : max_history_); @@ -85,7 +91,7 @@ class Batch : public Estimator { void Initialize(const mjModel* model) override; // reset memory - void Reset() override; + void Reset(const mjData* data = nullptr) override; // update void Update(const double* ctrl, const double* sensor) override; @@ -138,6 +144,22 @@ class Batch : public Estimator { mj_integratePos(model, q0, state + nq, -1.0 * model->opt.timestep); }; + // set time + void SetTime(double time) override { + // copy + double time_copy = time; + + // t1 + times.Set(&time_copy, 1); + + // t0 + time_copy -= model->opt.timestep; + times.Set(&time_copy, 0); + + // reset current time index + current_time_index = 1; + } + // set covariance void SetCovariance(const double* covariance) override { mju_copy(this->covariance.data(), covariance, ndstate_ * ndstate_); @@ -145,8 +167,10 @@ class Batch : public Estimator { }; // estimator-specific GUI elements - void GUI(mjUI& ui, double* process_noise, double* sensor_noise, - double& timestep, int& integrator) override; + void GUI(mjUI& ui, EstimatorGUIData& data) override; + + // set GUI data + void SetGUIData(EstimatorGUIData& data) override; // estimator-specific plots void Plots(mjvFigure* fig_planner, mjvFigure* fig_timer, int planner_shift, @@ -289,8 +313,9 @@ class Batch : public Estimator { bool assemble_sensor_norm_hessian = false; // assemble dense sensor norm Hessian bool assemble_force_norm_hessian = - false; // assemble dense force norm Hessian - bool filter = false; // filter mode + false; // assemble dense force norm Hessian + bool filter = false; // filter mode + bool recursive_prior_update = false; // recursively update prior matrix } settings; // finite-difference settings @@ -299,6 +324,9 @@ class Batch : public Estimator { bool flg_actuation = 1; } finite_difference; + // filter mode status + int current_time_index; + private: // convert sequence of configurations to velocities, accelerations void ConfigurationToVelocityAcceleration(); @@ -365,6 +393,12 @@ class Batch : public Estimator { const EstimatorTrajectory& configuration, const double* search_direction, double step_size); + // initialize filter mode + void InitializeFilter(); + + // shift head and resize trajectories + void ShiftResizeTrajectory(int new_head, int new_length); + // reset timers void ResetTimers(); @@ -401,6 +435,8 @@ class Batch : public Estimator { double cost_initial_; double cost_previous_; + // TODO(taylor): underscore + // cost gradient std::vector cost_gradient; // nv * max_history_ @@ -547,6 +583,14 @@ class Batch : public Estimator { std::vector scratch1_covariance_; // (nv * max_history_) * (nv * max_history_) + // conditioned matrix + std::vector mat00_; + std::vector mat10_; + std::vector mat11_; + std::vector condmat_; + std::vector scratch0_condmat_; + std::vector scratch1_condmat_; + // status (internal) int cost_count_; // number of cost evaluations bool cost_skip_ = false; // flag for only evaluating cost derivatives @@ -604,6 +648,20 @@ class Batch : public Estimator { // max history int max_history_ = 3; + + // trajectory cache + EstimatorTrajectory configuration_cache_; // nq x T + EstimatorTrajectory configuration_previous_cache_; // nq x T + EstimatorTrajectory velocity_cache_; // nv x T + EstimatorTrajectory acceleration_cache_; // nv x T + EstimatorTrajectory act_cache_; // na x T + EstimatorTrajectory times_cache_; // 1 x T + EstimatorTrajectory ctrl_cache_; // nu x T + EstimatorTrajectory sensor_measurement_cache_; // ns x T + EstimatorTrajectory sensor_prediction_cache_; // ns x T + EstimatorTrajectory sensor_mask_cache_; // num_sensor x T + EstimatorTrajectory force_measurement_cache_; // nv x T + EstimatorTrajectory force_prediction_cache_; // nv x T }; // estimator status string diff --git a/mjpc/estimators/estimator.h b/mjpc/estimators/estimator.h index 8956ee1f1..014427fed 100644 --- a/mjpc/estimators/estimator.h +++ b/mjpc/estimators/estimator.h @@ -20,6 +20,9 @@ #include #include +#include "mjpc/estimators/gui.h" +#include "mjpc/utilities.h" + namespace mjpc { // maximum terms @@ -36,47 +39,52 @@ class Estimator { virtual void Initialize(const mjModel* model) = 0; // reset memory - virtual void Reset() = 0; + virtual void Reset(const mjData* data = nullptr) = 0; - // update + // update virtual void Update(const double* ctrl, const double* sensor) = 0; - // get state + // get state virtual double* State() = 0; - // get covariance + // get covariance virtual double* Covariance() = 0; - // get time + // get time virtual double& Time() = 0; - // get model + // get model virtual mjModel* Model() = 0; - // get data + // get data virtual mjData* Data() = 0; - // process noise + // process noise virtual double* ProcessNoise() = 0; - // sensor noise + // sensor noise virtual double* SensorNoise() = 0; - // process dimension + // process dimension virtual int DimensionProcess() const = 0; - + // sensor dimension virtual int DimensionSensor() const = 0; // set state virtual void SetState(const double* state) = 0; + // set time + virtual void SetTime(double time) = 0; + // set covariance virtual void SetCovariance(const double* covariance) = 0; // estimator-specific GUI elements - virtual void GUI(mjUI& ui, double* process_noise, double* sensor_noise, - double& timestep, int& integrator) = 0; + virtual void GUI(mjUI& ui, EstimatorGUIData& data) = 0; + + // set GUI data + virtual void SetGUIData(EstimatorGUIData& data) = 0; // estimator-specific plots virtual void Plots(mjvFigure* fig_planner, mjvFigure* fig_timer, @@ -84,6 +92,185 @@ class Estimator { int* shift) = 0; }; +// ground truth estimator +class GroundTruth : public Estimator { + public: + // constructor + GroundTruth() = default; + GroundTruth(const mjModel* model) { + Initialize(model); + Reset(); + } + + // destructor + ~GroundTruth() { + if (data_) mj_deleteData(data_); + if (model) mj_deleteModel(model); + } + + // initialize + void Initialize(const mjModel* model) override { + // model + if (this->model) mj_deleteModel(this->model); + this->model = mj_copyModel(nullptr, model); + + // data + data_ = mj_makeData(model); + + // timestep + this->model->opt.timestep = GetNumberOrDefault(this->model->opt.timestep, + model, "estimator_timestep"); + + // -- dimensions -- // + ndstate_ = 2 * model->nv + model->na; + + // sensor start index + int sensor_start_ = GetNumberOrDefault(0, model, "estimator_sensor_start"); + + // number of sensors + int nsensor_ = + GetNumberOrDefault(model->nsensor, model, "estimator_number_sensor"); + + // sensor dimension + nsensordata_ = 0; + for (int i = 0; i < nsensor_; i++) { + nsensordata_ += model->sensor_dim[sensor_start_ + i]; + } + + // state + state.resize(model->nq + model->nv + model->na); + + // covariance + covariance.resize(ndstate_ * ndstate_); + + // process noise + noise_process.resize(ndstate_); + + // sensor noise + noise_sensor.resize(nsensordata_); // over allocate + } + + // reset + void Reset(const mjData* data = nullptr) override { + // dimensions + int nq = model->nq, nv = model->nv, na = model->na; + int ndstate = 2 * nv + na; + + if (data) { + mju_copy(state.data(), data->qpos, nq); + mju_copy(state.data() + nq, data->qvel, nv); + mju_copy(state.data() + nq + nv, data->act, na); + time = data->time; + } else { // model default + // set home keyframe + int home_id = mj_name2id(model, mjOBJ_KEY, "home"); + if (home_id >= 0) mj_resetDataKeyframe(model, data_, home_id); + + // state + mju_copy(state.data(), data_->qpos, nq); + mju_copy(state.data() + nq, data_->qvel, nv); + mju_copy(state.data() + nq + nv, data_->act, na); + time = data_->time; + } + + // covariance + mju_eye(covariance.data(), ndstate); + double covariance_scl = + GetNumberOrDefault(1.0e-4, model, "estimator_covariance_initial_scale"); + mju_scl(covariance.data(), covariance.data(), covariance_scl, + ndstate * ndstate); + + // process noise + double noise_process_scl = + GetNumberOrDefault(1.0e-4, model, "estimator_process_noise_scale"); + std::fill(noise_process.begin(), noise_process.end(), noise_process_scl); + + // sensor noise + double noise_sensor_scl = + GetNumberOrDefault(1.0e-4, model, "estimator_sensor_noise_scale"); + std::fill(noise_sensor.begin(), noise_sensor.end(), noise_sensor_scl); + } + + // update + void Update(const double* ctrl, const double* sensor) override {}; + + // get state + double* State() override { return state.data(); }; + + // get covariance + double* Covariance() override { return covariance.data(); }; + + // get time + double& Time() override { return time; }; + + // get model + mjModel* Model() override { return model; }; + + // get data + mjData* Data() override { return data_; }; + + // get process noise + double* ProcessNoise() override { return noise_process.data(); }; + + // get sensor noise + double* SensorNoise() override { return noise_sensor.data(); }; + + // process dimension + int DimensionProcess() const override { return ndstate_; }; + + // sensor dimension + int DimensionSensor() const override { return nsensordata_; }; + + // set state + void SetState(const double* state) override { + mju_copy(this->state.data(), state, ndstate_); + }; + + // set time + void SetTime(double time) override { + this->time = time; + } + + // set covariance + void SetCovariance(const double* covariance) override { + mju_copy(this->covariance.data(), covariance, ndstate_ * ndstate_); + }; + + // estimator-specific GUI elements + void GUI(mjUI& ui, EstimatorGUIData& data) override {}; + + // set GUI data + void SetGUIData(EstimatorGUIData& data) override {}; + + // estimator-specific plots + void Plots(mjvFigure* fig_planner, mjvFigure* fig_timer, int planner_shift, + int timer_shift, int planning, int* shift) override {}; + + // model + mjModel* model = nullptr; + + // state (nstate_) + std::vector state; + double time; + + // covariance (ndstate_ x ndstate_) + std::vector covariance; + + // process noise (ndstate_) + std::vector noise_process; + + // sensor noise (nsensordata_) + std::vector noise_sensor; + + private: + // data + mjData* data_ = nullptr; + + // dimensions + int ndstate_; + int nsensordata_; +}; + } // namespace mjpc -#endif // MJPC_ESTIMATORS_ESTIMATOR_H_ +#endif // MJPC_ESTIMATORS_ESTIMATOR_H_ \ No newline at end of file diff --git a/mjpc/estimators/gui.cc b/mjpc/estimators/gui.cc new file mode 100644 index 000000000..1f7bec4d5 --- /dev/null +++ b/mjpc/estimators/gui.cc @@ -0,0 +1,44 @@ +// Copyright 2022 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "mjpc/estimators/gui.h" + +#include + +#include + +#include "mjpc/utilities.h" + +namespace mjpc { + +// Initialize +void EstimatorGUIData::Initialize(const mjModel* model, int nprocess, + int nsensor) { + // timestep + timestep = model->opt.timestep; + + // integrator + integrator = model->opt.integrator; + + // process noise + process_noise.resize(nprocess); + + // sensor noise + sensor_noise.resize(nsensor); + + // estimation horizon + horizon = GetNumberOrDefault(3, model, "batch_configuration_length"); +} + +} // namespace mjpc \ No newline at end of file diff --git a/mjpc/estimators/gui.h b/mjpc/estimators/gui.h new file mode 100644 index 000000000..6ff82db43 --- /dev/null +++ b/mjpc/estimators/gui.h @@ -0,0 +1,54 @@ +// Copyright 2022 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJPC_ESTIMATORS_GUI_H_ +#define MJPC_ESTIMATORS_GUI_H_ + +#include + +#include + +namespace mjpc { + +// data that is modified by the GUI and later set in Estimator +class EstimatorGUIData { + public: + // constructor + EstimatorGUIData() = default; + + // destructor + ~EstimatorGUIData() = default; + + // Initialize + void Initialize(const mjModel* model, int nprocess, int nsensor); + + // time step + double timestep; + + // integrator + int integrator; + + // process noise + std::vector process_noise; + + // sensor noise + std::vector sensor_noise; + + // estimation horizon + int horizon; +}; + +} // namespace mjpc + +#endif // MJPC_ESTIMATORS_GUI_H_ \ No newline at end of file diff --git a/mjpc/estimators/include.cc b/mjpc/estimators/include.cc index 8a5a4c534..149336d3e 100644 --- a/mjpc/estimators/include.cc +++ b/mjpc/estimators/include.cc @@ -17,6 +17,7 @@ namespace mjpc { const char kEstimatorNames[] = + "Kalman\n" "Batch"; // load all available estimators @@ -25,7 +26,8 @@ std::vector> LoadEstimators() { std::vector> estimators; // add estimators - estimators.emplace_back(new mjpc::Batch(1)); // filter mode + estimators.emplace_back(new mjpc::GroundTruth()); // ground truth state + estimators.emplace_back(new mjpc::Kalman()); // extended Kalman filter return estimators; } diff --git a/mjpc/estimators/include.h b/mjpc/estimators/include.h index ec7c58cbf..c45050d72 100644 --- a/mjpc/estimators/include.h +++ b/mjpc/estimators/include.h @@ -20,6 +20,7 @@ #include "mjpc/estimators/batch.h" #include "mjpc/estimators/estimator.h" +#include "mjpc/estimators/kalman.h" namespace mjpc { diff --git a/mjpc/estimators/kalman.cc b/mjpc/estimators/kalman.cc new file mode 100644 index 000000000..57889ebf6 --- /dev/null +++ b/mjpc/estimators/kalman.cc @@ -0,0 +1,571 @@ +// Copyright 2023 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "mjpc/estimators/kalman.h" + +#include + +#include +#include + +#include "mjpc/array_safety.h" +#include "mjpc/estimators/estimator.h" +#include "mjpc/estimators/gui.h" +#include "mjpc/utilities.h" + +namespace mjpc { +namespace mju = ::mujoco::util_mjpc; + +// initialize +void Kalman::Initialize(const mjModel* model) { + // model + if (this->model) mj_deleteModel(this->model); + this->model = mj_copyModel(nullptr, model); + + // data + if (this->data_) mj_deleteData(this->data_); + data_ = mj_makeData(model); + + // timestep + this->model->opt.timestep = GetNumberOrDefault(this->model->opt.timestep, + model, "estimator_timestep"); + + // dimension + nstate_ = model->nq + model->nv + model->na; + ndstate_ = 2 * model->nv + model->na; + + // sensor start index + sensor_start_ = + GetNumberOrDefault(0, model, "estimator_sensor_start"); + + // number of sensors + nsensor_ = + GetNumberOrDefault(model->nsensor, model, "estimator_number_sensor"); + + // sensor dimension + nsensordata_ = 0; + for (int i = 0; i < nsensor_; i++) { + nsensordata_ += model->sensor_dim[sensor_start_ + i]; + } + + // sensor start index + sensor_start_index_ = 0; + for (int i = 0; i < sensor_start_; i++) { + sensor_start_index_ += model->sensor_dim[i]; + } + + // state + state.resize(nstate_); + + // covariance + covariance.resize(ndstate_ * ndstate_); + + // process noise + noise_process.resize(ndstate_); + + // sensor noise + noise_sensor.resize(nsensordata_); + + // dynamics Jacobian + dynamics_jacobian_.resize(ndstate_ * ndstate_); + + // sensor Jacobian + sensor_jacobian_.resize(model->nsensordata * ndstate_); + + // sensor error + sensor_error_.resize(nsensordata_); + + // correction + correction_.resize(ndstate_); + + // scratch + tmp0_.resize(ndstate_ * nsensordata_); + tmp1_.resize(nsensordata_ * nsensordata_); + tmp2_.resize(nsensordata_ * ndstate_); + tmp3_.resize(ndstate_ * ndstate_); +} + +// reset memory +void Kalman::Reset(const mjData* data) { + // dimension + int nq = model->nq, nv = model->nv, na = model->na; + + if (data) { + // state + mju_copy(state.data(), data->qpos, nq); + mju_copy(state.data() + nq, data->qvel, nv); + mju_copy(state.data() + nq + nv, data->act, na); + time = data->time; + } else { + // set home keyframe + int home_id = mj_name2id(model, mjOBJ_KEY, "home"); + if (home_id >= 0) mj_resetDataKeyframe(model, data_, home_id); + + // state + mju_copy(state.data(), data_->qpos, nq); + mju_copy(state.data() + nq, data_->qvel, nv); + mju_copy(state.data() + nq + nv, data_->act, na); + time = data_->time; + } + + // covariance + mju_eye(covariance.data(), ndstate_); + double covariance_scl = + GetNumberOrDefault(1.0e-4, model, "estimator_covariance_initial_scale"); + mju_scl(covariance.data(), covariance.data(), covariance_scl, + ndstate_ * ndstate_); + + // process noise + double noise_process_scl = + GetNumberOrDefault(1.0e-4, model, "estimator_process_noise_scale"); + std::fill(noise_process.begin(), noise_process.end(), noise_process_scl); + + // sensor noise + double noise_sensor_scl = + GetNumberOrDefault(1.0e-4, model, "estimator_sensor_noise_scale"); + std::fill(noise_sensor.begin(), noise_sensor.end(), noise_sensor_scl); + + // dynamics Jacobian + mju_zero(dynamics_jacobian_.data(), ndstate_ * ndstate_); + + // sensor Jacobian + mju_zero(sensor_jacobian_.data(), model->nsensordata * ndstate_); + + // sensor error + mju_zero(sensor_error_.data(), nsensordata_); + + // correction + mju_zero(correction_.data(), ndstate_); + + // timer + timer_measurement_ = 0.0; + timer_prediction_ = 0.0; + + // scratch + std::fill(tmp0_.begin(), tmp0_.end(), 0.0); + std::fill(tmp1_.begin(), tmp1_.end(), 0.0); + std::fill(tmp2_.begin(), tmp2_.end(), 0.0); + std::fill(tmp3_.begin(), tmp3_.end(), 0.0); +} + +// update measurement +void Kalman::UpdateMeasurement(const double* ctrl, const double* sensor) { + // start timer + auto start = std::chrono::steady_clock::now(); + + // dimensions + int nq = model->nq, nv = model->nv, na = model->na, nu = model->nu; + + // set state + mju_copy(data_->qpos, state.data(), nq); + mju_copy(data_->qvel, state.data() + nq, nv); + mju_copy(data_->act, state.data() + nq + nv, na); + + // set ctrl + mju_copy(data_->ctrl, ctrl, nu); + + // forward to get sensor + mj_forward(model, data_); + + mju_sub(sensor_error_.data(), sensor + sensor_start_index_, + data_->sensordata + sensor_start_index_, nsensordata_); + + // -- Kalman gain: P * C' (C * P * C' + R)^-1 -- // + + // sensor Jacobian + mjd_transitionFD(model, data_, settings.epsilon, settings.flg_centered, NULL, + NULL, sensor_jacobian_.data(), NULL); + + // grab rows + double* C = sensor_jacobian_.data() + sensor_start_index_ * ndstate_; + + // P * C' = tmp0 + mju_mulMatMatT(tmp0_.data(), covariance.data(), C, ndstate_, ndstate_, + nsensordata_); + + // C * P * C' = C * tmp0 = tmp1 + mju_mulMatMat(tmp1_.data(), C, tmp0_.data(), nsensordata_, ndstate_, + nsensordata_); + + // C * P * C' + R + for (int i = 0; i < nsensordata_; i++) { + tmp1_[nsensordata_ * i + i] += noise_sensor[i]; + } + + // factorize: C * P * C' + R + int rank = mju_cholFactor(tmp1_.data(), nsensordata_, 0.0); + if (rank < nsensordata_) { + // TODO(taylor): remove and return status + mju_error("measurement update rank: (%i / %i)\n", rank, nsensordata_); + } + + // -- correction: (P * C') * (C * P * C' + R)^-1 * sensor_error -- // + + // tmp2 = (C * P * C' + R) \ sensor_error + mju_cholSolve(tmp2_.data(), tmp1_.data(), sensor_error_.data(), nsensordata_); + + // correction = (P * C') * (C * P * C' + R) \ sensor_error = tmp0 * tmp2 + mju_mulMatVec(correction_.data(), tmp0_.data(), tmp2_.data(), ndstate_, + nsensordata_); + + // -- state update -- // + + // configuration + mj_integratePos(model, state.data(), correction_.data(), 1.0); + + // velocity + act + mju_addTo(state.data() + nq, correction_.data() + nv, nv + na); + + // -- covariance update -- // + // TODO(taylor): Joseph form update ? + + // tmp2 = (C * P * C' + R)^-1 (C * P) = tmp1 \ tmp0' + for (int i = 0; i < ndstate_; i++) { + mju_cholSolve(tmp2_.data() + nsensordata_ * i, tmp1_.data(), + tmp0_.data() + nsensordata_ * i, nsensordata_); + } + + // tmp3 = (P * C') * (C * P * C' + R)^-1 (C * P) = tmp0 * tmp2' + mju_mulMatMatT(tmp3_.data(), tmp0_.data(), tmp2_.data(), ndstate_, + nsensordata_, ndstate_); + + // covariance -= tmp3 + mju_subFrom(covariance.data(), tmp3_.data(), ndstate_ * ndstate_); + + // symmetrize + mju_symmetrize(covariance.data(), covariance.data(), ndstate_); + + // stop timer (ms) + timer_measurement_ = 1.0e-3 * GetDuration(start); +} + +// update time +void Kalman::UpdatePrediction() { + // start timer + auto start = std::chrono::steady_clock::now(); + + // dimensions + int nq = model->nq, nv = model->nv, na = model->na; + + // set state + mju_copy(data_->qpos, state.data(), nq); + mju_copy(data_->qvel, state.data() + nq, nv); + mju_copy(data_->act, state.data() + nq + nv, na); + + // dynamics Jacobian + mjd_transitionFD(model, data_, settings.epsilon, settings.flg_centered, + dynamics_jacobian_.data(), NULL, NULL, NULL); + + // integrate state + mj_step(model, data_); + + // update state + mju_copy(state.data(), data_->qpos, nq); + mju_copy(state.data() + nq, data_->qvel, nv); + mju_copy(state.data() + nq + nv, data_->act, na); + + // -- update covariance: P = A * P * A' -- // + + // tmp = P * A' + mju_mulMatMatT(tmp3_.data(), covariance.data(), dynamics_jacobian_.data(), + ndstate_, ndstate_, ndstate_); + + // P = A * tmp + mju_mulMatMat(covariance.data(), dynamics_jacobian_.data(), tmp3_.data(), + ndstate_, ndstate_, ndstate_); + + // process noise + for (int i = 0; i < ndstate_; i++) { + covariance[ndstate_ * i + i] += noise_process[i]; + } + + // symmetrize + mju_symmetrize(covariance.data(), covariance.data(), ndstate_); + + // stop timer + timer_prediction_ = 1.0e-3 * GetDuration(start); +} + +// estimator-specific GUI elements +void Kalman::GUI(mjUI& ui, EstimatorGUIData& data) { + + // ----- estimator ------ // + mjuiDef defEstimator[] = { + {mjITEM_SECTION, "Estimator Settings", 1, nullptr, + "AP"}, // needs new section to satisfy mjMAXUIITEM + {mjITEM_BUTTON, "Reset", 2, nullptr, ""}, + {mjITEM_SLIDERNUM, "Timestep", 2, &data.timestep, "1.0e-3 0.1"}, + {mjITEM_SELECT, "Integrator", 2, &data.integrator, + "Euler\nRK4\nImplicit\nFastImplicit"}, + {mjITEM_END}}; + + // add estimator + mjui_add(&ui, defEstimator); + + // -- process noise -- // + int nv = model->nv; + int process_noise_shift = 0; + mjuiDef defProcessNoise[kMaxProcessNoise + 2]; + + // separator + defProcessNoise[0] = {mjITEM_SEPARATOR, "Process Noise Covariance", 1}; + process_noise_shift++; + + // add UI elements + for (int i = 0; i < DimensionProcess(); i++) { + // element + defProcessNoise[process_noise_shift] = {mjITEM_SLIDERNUM, "", 2, + data.process_noise.data() + i, "1.0e-8 0.01"}; + + // set name + mju::strcpy_arr(defProcessNoise[process_noise_shift].name, ""); + + // shift + process_noise_shift++; + } + + // name UI elements + int jnt_shift = 1; + std::string jnt_name_pos; + std::string jnt_name_vel; + + // loop over joints + for (int i = 0; i < model->njnt; i++) { + int name_jntadr = model->name_jntadr[i]; + std::string jnt_name(model->names + name_jntadr); + + // get joint type + int jnt_type = model->jnt_type[i]; + + // free + switch (jnt_type) { + case mjJNT_FREE: + // position + jnt_name_pos = jnt_name + " (pos 0)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 0].name, + jnt_name_pos.c_str()); + + jnt_name_pos = jnt_name + " (pos 1)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 1].name, + jnt_name_pos.c_str()); + + jnt_name_pos = jnt_name + " (pos 2)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 2].name, + jnt_name_pos.c_str()); + + jnt_name_pos = jnt_name + " (pos 3)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 3].name, + jnt_name_pos.c_str()); + + jnt_name_pos = jnt_name + " (pos 4)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 4].name, + jnt_name_pos.c_str()); + + jnt_name_pos = jnt_name + " (pos 5)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 5].name, + jnt_name_pos.c_str()); + + // velocity + jnt_name_vel = jnt_name + " (vel 0)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 0].name, + jnt_name_vel.c_str()); + + jnt_name_vel = jnt_name + " (vel 1)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 1].name, + jnt_name_vel.c_str()); + + jnt_name_vel = jnt_name + " (vel 2)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 2].name, + jnt_name_vel.c_str()); + + jnt_name_vel = jnt_name + " (vel 3)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 3].name, + jnt_name_vel.c_str()); + + jnt_name_vel = jnt_name + " (vel 4)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 4].name, + jnt_name_vel.c_str()); + + jnt_name_vel = jnt_name + " (vel 5)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 5].name, + jnt_name_vel.c_str()); + + // shift + jnt_shift += 6; + break; + case mjJNT_BALL: + // position + jnt_name_pos = jnt_name + " (pos 0)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 0].name, + jnt_name_pos.c_str()); + + jnt_name_pos = jnt_name + " (pos 1)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 1].name, + jnt_name_pos.c_str()); + + jnt_name_pos = jnt_name + " (pos 2)"; + mju::strcpy_arr(defProcessNoise[jnt_shift + 2].name, + jnt_name_pos.c_str()); + + // velocity + jnt_name_vel = jnt_name + " (vel 0)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 0].name, + jnt_name_vel.c_str()); + + jnt_name_vel = jnt_name + " (vel 1)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 1].name, + jnt_name_vel.c_str()); + + jnt_name_vel = jnt_name + " (vel 2)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + 2].name, + jnt_name_vel.c_str()); + + // shift + jnt_shift += 3; + break; + case mjJNT_HINGE: + // position + jnt_name_pos = jnt_name + " (pos)"; + mju::strcpy_arr(defProcessNoise[jnt_shift].name, jnt_name_pos.c_str()); + + // velocity + jnt_name_vel = jnt_name + " (vel)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift].name, + jnt_name_vel.c_str()); + + // shift + jnt_shift++; + break; + case mjJNT_SLIDE: + // position + jnt_name_pos = jnt_name + " (pos)"; + mju::strcpy_arr(defProcessNoise[jnt_shift].name, jnt_name_pos.c_str()); + + // velocity + jnt_name_vel = jnt_name + " (vel)"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift].name, + jnt_name_vel.c_str()); + + // shift + jnt_shift++; + break; + } + } + + // loop over act + std::string act_str; + for (int i = 0; i < model->na; i++) { + act_str = "act (" + std::to_string(i) + ")"; + mju::strcpy_arr(defProcessNoise[nv + jnt_shift + i].name, act_str.c_str()); + } + + // end + defProcessNoise[process_noise_shift] = {mjITEM_END}; + + // add process noise + mjui_add(&ui, defProcessNoise); + + // -- sensor noise -- // + int sensor_noise_shift = 0; + mjuiDef defSensorNoise[kMaxSensorNoise + 2]; + + // separator + defSensorNoise[0] = {mjITEM_SEPARATOR, "Sensor Noise Covariance", 1}; + sensor_noise_shift++; + + // loop over sensors + std::string sensor_str; + for (int i = 0; i < nsensor_; i++) { + std::string name_sensor(model->names + + model->name_sensoradr[sensor_start_ + i]); + int dim_sensor = model->sensor_dim[sensor_start_ + i]; + + // loop over sensor elements + for (int j = 0; j < dim_sensor; j++) { + // element + defSensorNoise[sensor_noise_shift] = { + mjITEM_SLIDERNUM, "", 2, data.sensor_noise.data() + sensor_noise_shift - 1, + "1.0e-8 0.01"}; + + // sensor name + sensor_str = name_sensor; + + // add element index + if (dim_sensor > 1) { + sensor_str += " (" + std::to_string(j) + ")"; + } + + // set sensor name + mju::strcpy_arr(defSensorNoise[sensor_noise_shift].name, + sensor_str.c_str()); + + // shift + sensor_noise_shift++; + } + } + + // end + defSensorNoise[sensor_noise_shift] = {mjITEM_END}; + + // add sensor noise + mjui_add(&ui, defSensorNoise); +} + +// set GUI data +void Kalman::SetGUIData(EstimatorGUIData& data) { + mju_copy(noise_process.data(), + data.process_noise.data(), + DimensionProcess()); + mju_copy(noise_sensor.data(), + data.sensor_noise.data(), + DimensionSensor()); + model->opt.timestep = data.timestep; + model->opt.integrator = data.integrator; +}; + +// estimator-specific plots +void Kalman::Plots(mjvFigure* fig_planner, mjvFigure* fig_timer, int planner_shift, + int timer_shift, int planning, int* shift) { + // Kalman info + double estimator_bounds[2] = {-6, 6}; + + // covariance trace + double trace = Trace(covariance.data(), DimensionProcess()); + mjpc::PlotUpdateData(fig_planner, estimator_bounds, + fig_planner->linedata[planner_shift + 0][0] + 1, + mju_log10(trace), 100, planner_shift + 0, 0, 1, -100); + + // legend + mju::strcpy_arr(fig_planner->linename[planner_shift + 0], + "Covariance Trace"); + + // Kalman timers + double timer_bounds[2] = {0.0, 1.0}; + + // measurement update + PlotUpdateData(fig_timer, timer_bounds, + fig_timer->linedata[timer_shift + 0][0] + 1, + TimerMeasurement(), 100, timer_shift + 0, 0, 1, -100); + + // prediction update + PlotUpdateData(fig_timer, timer_bounds, + fig_timer->linedata[timer_shift + 1][0] + 1, TimerPrediction(), + 100, timer_shift + 1, 0, 1, -100); + + // legend + mju::strcpy_arr(fig_timer->linename[timer_shift + 0], "Measurement Update"); + mju::strcpy_arr(fig_timer->linename[timer_shift + 1], "Prediction Update"); +} + +} // namespace mjpc \ No newline at end of file diff --git a/mjpc/estimators/kalman.h b/mjpc/estimators/kalman.h new file mode 100644 index 000000000..6df1840ed --- /dev/null +++ b/mjpc/estimators/kalman.h @@ -0,0 +1,188 @@ +// Copyright 2023 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MJPC_ESTIMATORS_KALMAN_H_ +#define MJPC_ESTIMATORS_KALMAN_H_ + +#include + +#include +#include + +#include "mjpc/estimators/estimator.h" +#include "mjpc/estimators/gui.h" +#include "mjpc/utilities.h" + +namespace mjpc { + +// https://stanford.edu/class/ee363/lectures/kf.pdf +class Kalman : public Estimator { + public: + // constructor + Kalman() = default; + Kalman(const mjModel* model) { + Initialize(model); + Reset(); + } + + // destructor + ~Kalman() { + if (data_) mj_deleteData(data_); + if (model) mj_deleteModel(model); + } + + // initialize + void Initialize(const mjModel* model) override; + + // reset memory + void Reset(const mjData* data = nullptr) override; + + // update measurement + void UpdateMeasurement(const double* ctrl, const double* sensor); + + // update time + void UpdatePrediction(); + + // update + void Update(const double* ctrl, const double* sensor) override { + // correct state with latest measurement + UpdateMeasurement(ctrl, sensor); + + // propagate state forward in time with model + UpdatePrediction(); + + // set time + time = data_->time; + } + + // get state + double* State() override { return state.data(); }; + + // get covariance + double* Covariance() override { return covariance.data(); }; + + // get time + double& Time() override { return time; }; + + // get model + mjModel* Model() override { return model; }; + + // get data + mjData* Data() override { return data_; }; + + // get process noise + double* ProcessNoise() override { return noise_process.data(); }; + + // get sensor noise + double* SensorNoise() override { return noise_sensor.data(); }; + + // dimension process + int DimensionProcess() const override { return ndstate_; }; + + // dimension sensor + int DimensionSensor() const override { return nsensordata_; }; + + // set state + void SetState(const double* state) override { + mju_copy(this->state.data(), state, ndstate_); + }; + + // set time + void SetTime(double time) override { + this->time = time; + } + + // set covariance + void SetCovariance(const double* covariance) override { + mju_copy(this->covariance.data(), covariance, ndstate_ * ndstate_); + }; + + // get measurement timer (ms) + double TimerMeasurement() const { return timer_measurement_; }; + + // get prediction timer (ms) + double TimerPrediction() const { return timer_prediction_; }; + + // estimator-specific GUI elements + void GUI(mjUI& ui, EstimatorGUIData& data) override; + + // set GUI data + void SetGUIData(EstimatorGUIData& data) override; + + // estimator-specific plots + void Plots(mjvFigure* fig_planner, mjvFigure* fig_timer, int planner_shift, + int timer_shift, int planning, int* shift) override; + + // model + mjModel* model = nullptr; + + // state (nstate_) + std::vector state; + double time; + + // covariance (ndstate_ x ndstate_) + std::vector covariance; + + // process noise (ndstate_) + std::vector noise_process; + + // sensor noise (nsensordata_) + std::vector noise_sensor; + + // settings + struct Settings { + double epsilon = 1.0e-6; + bool flg_centered = false; + } settings; + + private: + // dimensions + int nstate_; + int ndstate_; + int nsensordata_; + int nsensor_; + + // sensor indexing + int sensor_start_; + int sensor_start_index_; + + // data + mjData* data_ = nullptr; + + // correction (ndstate_) + std::vector correction_; + + // sensor Jacobian (nsensordata x ndstate_) + std::vector sensor_jacobian_; + + // dynamics Jacobian (ndstate_ x ndstate_) + std::vector dynamics_jacobian_; + + // sensor error (nsensordata_) + std::vector sensor_error_; + + // timer (ms) + double timer_measurement_; + double timer_prediction_; + + // scratch + std::vector tmp0_; + std::vector tmp1_; + std::vector tmp2_; + std::vector tmp3_; +}; + +} // namespace mjpc + +#endif // MJPC_ESTIMATORS_KALMAN_H_ \ No newline at end of file diff --git a/mjpc/test/CMakeLists.txt b/mjpc/test/CMakeLists.txt index 112dac539..8f5ba2337 100644 --- a/mjpc/test/CMakeLists.txt +++ b/mjpc/test/CMakeLists.txt @@ -64,12 +64,12 @@ target_include_directories(simulation PRIVATE ${mujoco_SOURCE_DIR}/include gmock enable_testing() -add_subdirectory(agent) +# add_subdirectory(agent) add_subdirectory(estimator) -add_subdirectory(gradient_planner) -add_subdirectory(ilqg_planner) -add_subdirectory(planners/robust) -add_subdirectory(sampling_planner) -add_subdirectory(state) -add_subdirectory(tasks) -add_subdirectory(utilities) +# add_subdirectory(gradient_planner) +# add_subdirectory(ilqg_planner) +# add_subdirectory(planners/robust) +# add_subdirectory(sampling_planner) +# add_subdirectory(state) +# add_subdirectory(tasks) +# add_subdirectory(utilities) diff --git a/mjpc/test/estimator/CMakeLists.txt b/mjpc/test/estimator/CMakeLists.txt index 5b58f03fe..b0a5ce67b 100644 --- a/mjpc/test/estimator/CMakeLists.txt +++ b/mjpc/test/estimator/CMakeLists.txt @@ -29,3 +29,6 @@ target_link_libraries(estimator_trajectory_test gmock) test(estimator_utilities_test) target_link_libraries(estimator_utilities_test load simulation gmock) + +test(kalman_test) +target_link_libraries(kalman_test load simulation gmock) \ No newline at end of file diff --git a/mjpc/test/estimator/batch_prior_test.cc b/mjpc/test/estimator/batch_prior_test.cc index 3d4d9bc9e..57fefb916 100644 --- a/mjpc/test/estimator/batch_prior_test.cc +++ b/mjpc/test/estimator/batch_prior_test.cc @@ -118,15 +118,15 @@ TEST(PriorCost, Particle) { // cost double cost_lambda = cost_prior(estimator.configuration.Data()); - // gradient + // // gradient FiniteDifferenceGradient fdg(nvar); fdg.Compute(cost_prior, estimator.configuration.Data(), nvar); - // Hessian + // // Hessian FiniteDifferenceHessian fdh(nvar); fdh.Compute(cost_prior, estimator.configuration.Data(), nvar); - // ----- estimator ----- // + // // ----- estimator ----- // ThreadPool pool(1); // evaluate cost diff --git a/mjpc/test/estimator/batch_sensor_test.cc b/mjpc/test/estimator/batch_sensor_test.cc index 00247a2c8..98ec5aa67 100644 --- a/mjpc/test/estimator/batch_sensor_test.cc +++ b/mjpc/test/estimator/batch_sensor_test.cc @@ -90,7 +90,7 @@ TEST(MeasurementCost, Particle) { // initialize double cost = 0.0; - // time scaling + // time scaling double time_scale = 1.0; double time_scale2 = 1.0; if (estimator.settings.time_scaling_sensor) { @@ -119,7 +119,7 @@ TEST(MeasurementCost, Particle) { // loop over sensors int shift = 0; - for (int i = 0; i < model->nsensor; i++) { + for (int i = 0; i < model->nsensor; i++) { // skip velocity, acceleration sensors (assumes position sensors are // first) if (model->sensor_needstage[i] != mjSTAGE_POS) continue; @@ -180,7 +180,7 @@ TEST(MeasurementCost, Particle) { int shift = 0; for (int i = 0; i < model->nsensor; i++) { - // sensor stage + // sensor stage int sensor_stage = model->sensor_needstage[i]; // sensor dimension @@ -189,7 +189,7 @@ TEST(MeasurementCost, Particle) { // sensor residual double* rki = rk + shift; - // time weight + // time weight double time_weight = 1.0; if (sensor_stage == mjSTAGE_VEL) { time_weight = time_scale; @@ -356,7 +356,7 @@ TEST(MeasurementCost, Box) { // initialize double cost = 0.0; - // time scale + // time scale double time_scale = 1.0; double time_scale2 = 1.0; if (estimator.settings.time_scaling_sensor) { @@ -446,7 +446,7 @@ TEST(MeasurementCost, Box) { int shift = 0; for (int i = 0; i < model->nsensor; i++) { - // sensor stage + // sensor stage int sensor_stage = model->sensor_needstage[i]; // sensor dimension @@ -455,7 +455,7 @@ TEST(MeasurementCost, Box) { // sensor residual double* rki = rk + shift; - // time weight + // time weight double time_weight = 1.0; if (sensor_stage == mjSTAGE_VEL) { time_weight = time_scale; diff --git a/mjpc/test/estimator/estimator_utilities_test.cc b/mjpc/test/estimator/estimator_utilities_test.cc index 6398b55b2..cc23244f0 100644 --- a/mjpc/test/estimator/estimator_utilities_test.cc +++ b/mjpc/test/estimator/estimator_utilities_test.cc @@ -48,7 +48,7 @@ TEST(FiniteDifferenceVelocityAcceleration, Particle2D) { }; // trajectories - int T = 200; + int T = 10; std::vector qpos(nq * T); std::vector qvel(nv * T); std::vector qacc(nv * T); diff --git a/mjpc/test/estimator/kalman_test.cc b/mjpc/test/estimator/kalman_test.cc new file mode 100644 index 000000000..4a37e1667 --- /dev/null +++ b/mjpc/test/estimator/kalman_test.cc @@ -0,0 +1,115 @@ +// Copyright 2023 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "mjpc/estimators/kalman.h" + +#include +#include + +#include + +#include "gtest/gtest.h" +#include "mjpc/estimators/trajectory.h" +#include "mjpc/test/load.h" +#include "mjpc/test/simulation.h" +#include "mjpc/utilities.h" + +namespace mjpc { +namespace { + +TEST(Estimator, Kalman) { + // load model + mjModel* model = LoadTestModel("estimator/particle/task_imu.xml"); + mjData* data = mj_makeData(model); + + // ----- rollout ----- // + int T = 200; + Simulation sim(model, T); + auto controller = [](double* ctrl, double time) {}; + double qpos0[1] = {0.25}; + sim.SetState(qpos0, NULL); + sim.Rollout(controller); + + // ----- Kalman ----- // + + // initialize filter + Kalman kalman(model); + + // set initial state + mju_copy(kalman.state.data(), sim.qpos.Get(0), model->nq); + mju_copy(kalman.state.data() + model->nq, sim.qvel.Get(0), model->nv); + + // initialize covariance + mju_eye(kalman.covariance.data(), 2 * model->nv); + mju_scl(kalman.covariance.data(), kalman.covariance.data(), 1.0e-5, + (2 * model->nv) * (2 * model->nv)); + + // initial process noise + mju_fill(kalman.noise_process.data(), 1.0e-5, 2 * model->nv); + + // initialize sensor noise + mju_fill(kalman.noise_sensor.data(), 1.0e-5, model->nsensordata); + + // Kalman trajectories + EstimatorTrajectory kalman_qpos(model->nq, T); + EstimatorTrajectory kalman_qvel(model->nv, T); + EstimatorTrajectory kalman_timer_measurement(1, T); + EstimatorTrajectory kalman_timer_prediction(1, T); + + // noisy sensor + std::vector noisy_sensor(model->nsensordata); + absl::BitGen gen_; + + for (int t = 0; t < T; t++) { + // noisy sensor + mju_copy(noisy_sensor.data(), sim.sensor.Get(t), model->nsensordata); + for (int i = 0; i < model->nsensordata; i++) { + noisy_sensor[i] += 0.0 * absl::Gaussian(gen_, 0.0, 1.0); + } + // measurement update + kalman.UpdateMeasurement(sim.ctrl.Get(t), noisy_sensor.data()); + + // cache state + kalman_qpos.Set(kalman.state.data(), t); + kalman_qvel.Set(kalman.state.data() + model->nq, t); + + // cache timer + double timer_measurement = kalman.TimerMeasurement(); + kalman_timer_measurement.Set(&timer_measurement, t); + + // prediction update + kalman.UpdatePrediction(); + + // cache timer + double timer_prediction = kalman.TimerPrediction(); + kalman_timer_prediction.Set(&timer_prediction, t); + + // test qpos + std::vector pos_error(model->nq); + mju_sub(pos_error.data(), kalman_qpos.Get(t), sim.qpos.Get(t), model->nq); + EXPECT_NEAR(mju_norm(pos_error.data(), model->nq), 0.0, 1.0e-4); + + // test qvel + std::vector vel_error(model->nv); + mju_sub(vel_error.data(), kalman_qvel.Get(t), sim.qvel.Get(t), model->nv); + EXPECT_NEAR(mju_norm(vel_error.data(), model->nv), 0.0, 1.0e-4); + } + + // delete data + model + mj_deleteData(data); + mj_deleteModel(model); +} + +} // namespace +} // namespace mjpc \ No newline at end of file diff --git a/mjpc/test/testdata/estimator/box/task2D.xml b/mjpc/test/testdata/estimator/box/task2D.xml index 428c63623..4c69682bd 100644 --- a/mjpc/test/testdata/estimator/box/task2D.xml +++ b/mjpc/test/testdata/estimator/box/task2D.xml @@ -46,6 +46,14 @@ + diff --git a/mjpc/test/testdata/estimator/box/task3Drot.xml b/mjpc/test/testdata/estimator/box/task3Drot.xml new file mode 100644 index 000000000..aab3b9ec3 --- /dev/null +++ b/mjpc/test/testdata/estimator/box/task3Drot.xml @@ -0,0 +1,65 @@ + + \ No newline at end of file diff --git a/mjpc/test/testdata/estimator/box/task3Drot2.xml b/mjpc/test/testdata/estimator/box/task3Drot2.xml new file mode 100644 index 000000000..59d2cdd70 --- /dev/null +++ b/mjpc/test/testdata/estimator/box/task3Drot2.xml @@ -0,0 +1,75 @@ + + \ No newline at end of file diff --git a/mjpc/test/testdata/estimator/particle/task_imu.xml b/mjpc/test/testdata/estimator/particle/task_imu.xml new file mode 100644 index 000000000..d0835cc64 --- /dev/null +++ b/mjpc/test/testdata/estimator/particle/task_imu.xml @@ -0,0 +1,56 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/mjpc/test/testdata/estimator/quadruped/task.xml b/mjpc/test/testdata/estimator/quadruped/task.xml index b0c8269d0..87bfa33d6 100644 --- a/mjpc/test/testdata/estimator/quadruped/task.xml +++ b/mjpc/test/testdata/estimator/quadruped/task.xml @@ -1,5 +1,4 @@ - @@ -31,10 +30,6 @@ - - - - diff --git a/python/mujoco_mpc/batch.py b/python/mujoco_mpc/batch.py index 36c310f2c..b2251bed9 100644 --- a/python/mujoco_mpc/batch.py +++ b/python/mujoco_mpc/batch.py @@ -486,7 +486,7 @@ def status_code(code): elif code == 3: return "SMALL_DIRECTION_FAILURE" elif code == 4: - return "MAX_REGULARIZATION_FAILURE" + return "kMaxBatchRegularization_FAILURE" elif code == 5: return "COST_DIFFERENCE_FAILURE" elif code == 6: diff --git a/python/mujoco_mpc/kalman.py b/python/mujoco_mpc/kalman.py new file mode 100644 index 000000000..4ff8ff7d1 --- /dev/null +++ b/python/mujoco_mpc/kalman.py @@ -0,0 +1,270 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Python interface for interface with Kalman.""" + +import atexit +import os +import pathlib +import socket +import subprocess +import sys +import tempfile +from typing import Literal, Optional + +import grpc +import mujoco +import numpy as np +from numpy import typing as npt + +# INTERNAL IMPORT +from mujoco_mpc.proto import kalman_pb2 +from mujoco_mpc.proto import kalman_pb2_grpc + + +def find_free_port() -> int: + """Find an available TCP port on the system. + + This function creates a temporary socket, binds it to an available port + chosen by the operating system, and returns the chosen port number. + + Returns: + int: An available TCP port number. + """ + with socket.socket(family=socket.AF_INET6) as s: + s.bind(("", 0)) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + return s.getsockname()[1] + + +class Kalman: + """`Kalman` class to interface with MuJoCo MPC kalman. + + Attributes: + port: + channel: + stub: + server_process: + """ + + def __init__( + self, + model: mujoco.MjModel, + server_binary_path: Optional[str] = None, + send_as: Literal["mjb", "xml"] = "xml", + colab_logging: bool = True, + ): + # server + if server_binary_path is None: + binary_name = "kalman_server" + server_binary_path = pathlib.Path(__file__).parent / "mjpc" / binary_name + self._colab_logging = colab_logging + self.port = find_free_port() + self.server_process = subprocess.Popen( + [str(server_binary_path), f"--mjpc_port={self.port}"], + stdout=subprocess.PIPE if colab_logging else None, + ) + os.set_blocking(self.server_process.stdout.fileno(), False) + atexit.register(self.server_process.kill) + + credentials = grpc.local_channel_credentials(grpc.LocalConnectionType.LOCAL_TCP) + self.channel = grpc.secure_channel(f"localhost:{self.port}", credentials) + grpc.channel_ready_future(self.channel).result(timeout=10) + self.stub = kalman_pb2_grpc.KalmanStub(self.channel) + + # initialize + self.init( + model, + send_as=send_as, + ) + + def close(self): + self.channel.close() + self.server_process.kill() + self.server_process.wait() + + def init( + self, + model: mujoco.MjModel, + send_as: Literal["mjb", "xml"] = "xml", + ): + """ + Args: + model: optional `MjModel` instance, which, if provided, will be used as + the underlying model for planning. If not provided, the default MJPC + task xml will be used. + configuration_length: estimation horizon. + send_as: The serialization format for sending the model over gRPC; "xml". + """ + + # setup model + def model_to_mjb(model: mujoco.MjModel) -> bytes: + buffer_size = mujoco.mj_sizeModel(model) + buffer = np.empty(shape=buffer_size, dtype=np.uint8) + mujoco.mj_saveModel(model, None, buffer) + return buffer.tobytes() + + def model_to_xml(model: mujoco.MjModel) -> str: + tmp = tempfile.NamedTemporaryFile() + mujoco.mj_saveLastXML(tmp.name, model) + with pathlib.Path(tmp.name).open("rt") as f: + xml_string = f.read() + return xml_string + + if model is not None: + if send_as == "mjb": + model_message = kalman_pb2.MjModel(mjb=model_to_mjb(model)) + else: + model_message = kalman_pb2.MjModel(xml=model_to_xml(model)) + else: + model_message = None + + # initialize request + init_request = kalman_pb2.InitRequest( + model=model_message, + ) + + # initialize response + self._wait(self.stub.Init.future(init_request)) + + def reset(self): + # reset request + request = kalman_pb2.ResetRequest() + + # reset response + self._wait(self.stub.Reset.future(request)) + + def settings( + self, + epsilon: Optional[float] = None, + flg_centered: Optional[bool] = None, + ) -> dict[str, int | bool]: + # assemble settings + inputs = kalman_pb2.Settings( + epsilon=epsilon, + flg_centered=flg_centered, + ) + + # settings request + request = kalman_pb2.SettingsRequest( + settings=inputs, + ) + + # settings response + settings = self._wait(self.stub.Settings.future(request)).settings + + # return all settings + return { + "epsilon": settings.epsilon, + "flg_centered": settings.flg_centered, + } + + def update_measurement( + self, ctrl: Optional[npt.ArrayLike] = [], sensor: Optional[npt.ArrayLike] = [] + ): + # request + request = kalman_pb2.UpdateMeasurementRequest( + ctrl=ctrl, + sensor=sensor, + ) + + # response + self._wait(self.stub.UpdateMeasurement.future(request)) + + def update_prediction(self): + # request + request = kalman_pb2.UpdatePredictionRequest() + + # response + self._wait(self.stub.UpdatePrediction.future(request)) + + def timers(self) -> dict[str, float]: + # request + request = kalman_pb2.TimersRequest() + + # response + response = self._wait(self.stub.Timers.future(request)) + + # timers + return { + "measurement": response.measurement, + "prediction": response.prediction, + } + + def state(self, state: Optional[npt.ArrayLike] = []) -> np.ndarray: + # input + input = kalman_pb2.State(state=state) + + # request + request = kalman_pb2.StateRequest( + state=input, + ) + + # response + response = self._wait(self.stub.State.future(request)).state + + # return state + return np.array(response.state) + + def covariance(self, covariance: Optional[npt.ArrayLike] = None) -> np.ndarray: + # input + inputs = kalman_pb2.Covariance( + covariance=covariance.flatten() if covariance is not None else None, + ) + + # request + request = kalman_pb2.CovarianceRequest( + covariance=inputs, + ) + + # response + response = self._wait(self.stub.Covariance.future(request)).covariance + + # return covariance + return np.array(response.covariance).reshape(response.dimension, response.dimension) + + def noise( + self, process: Optional[npt.ArrayLike] = [], sensor: Optional[npt.ArrayLike] = [] + ) -> dict[str, np.ndarray]: + # inputs + inputs = kalman_pb2.Noise( + process=process, + sensor=sensor, + ) + + # request + request = kalman_pb2.NoiseRequest( + noise=inputs, + ) + + # response + response = self._wait(self.stub.Noise.future(request)).noise + + # return noise + return { + "process": np.array(response.process), + "sensor": np.array(response.sensor), + } + + def _wait(self, future): + """Waits for the future to complete, while printing out subprocess stdout.""" + if self._colab_logging: + while True: + line = self.server_process.stdout.readline() + if line: + sys.stdout.write(line.decode("utf-8")) + if future.done(): + break + return future.result() \ No newline at end of file diff --git a/python/mujoco_mpc/kalman_test.py b/python/mujoco_mpc/kalman_test.py new file mode 100644 index 000000000..a66041037 --- /dev/null +++ b/python/mujoco_mpc/kalman_test.py @@ -0,0 +1,100 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +from absl.testing import absltest +import mujoco +from mujoco_mpc import kalman as kalman_lib +import numpy as np + +import pathlib + + +class KalmanTest(absltest.TestCase): + + def test_settings(self): + # load model + model_path = ( + pathlib.Path(__file__).parent.parent.parent + / "mjpc/test/testdata/estimator/particle/task1D.xml" + ) + model = mujoco.MjModel.from_xml_path(str(model_path)) + + # initialize + kalman = kalman_lib.Kalman(model=model) + + # settings + epsilon = 2.0 + flg_centered = True + settings = kalman.settings( + epsilon=epsilon, flg_centered=flg_centered + ) + + # test + self.assertLess(np.abs(settings["epsilon"] - epsilon), 1.0e-5) + self.assertTrue(settings["flg_centered"] == flg_centered) + + def test_updates(self): + # load model + model_path = ( + pathlib.Path(__file__).parent.parent.parent + / "mjpc/test/testdata/estimator/particle/task1D.xml" + ) + model = mujoco.MjModel.from_xml_path(str(model_path)) + + # initialize + kalman = kalman_lib.Kalman(model=model) + + # state + state = np.random.normal(scale=1.0, size=(model.nq + model.nv)) + state_response = kalman.state(state=state) + + # test state + self.assertLess(np.linalg.norm(state_response - state), 1.0e-5) + + # covariance + nvelocity = 2 * model.nv + F = np.random.normal(scale=1.0, size=(nvelocity**2)).reshape(nvelocity, nvelocity) + covariance = F.T @ F + covariance_response = kalman.covariance(covariance=covariance) + + # test covariance + self.assertLess(np.linalg.norm((covariance_response - covariance).ravel()), 1.0e-5) + self.assertTrue(covariance_response.shape == (nvelocity, nvelocity)) + + # noise + process = np.random.normal(scale=1.0e-3, size=nvelocity) + sensor = np.random.normal(scale=1.0e-3, size=model.nsensordata) + noise = kalman.noise(process=process, sensor=sensor) + + # test noise + self.assertLess(np.linalg.norm(noise["process"] - process), 1.0e-5) + self.assertLess(np.linalg.norm(noise["sensor"] - sensor), 1.0e-5) + + # measurement update + ctrl = np.random.normal(scale=1.0, size=model.nu) + sensor = np.random.normal(scale=1.0, size=model.nsensordata) + kalman.update_measurement(ctrl=ctrl, sensor=sensor) + + # # prediction update + kalman.update_prediction() + + # timers + timer = kalman.timers() + + self.assertTrue(timer["measurement"] > 0.0) + self.assertTrue(timer["prediction"] > 0.0) + +if __name__ == "__main__": + absltest.main() \ No newline at end of file diff --git a/python/setup_kalman.py b/python/setup_kalman.py new file mode 100644 index 000000000..58963564a --- /dev/null +++ b/python/setup_kalman.py @@ -0,0 +1,299 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Install script for MuJoCo MPC.""" + +import os +import pathlib +import platform +import shutil +import setuptools +from setuptools.command import build_ext +from setuptools.command import build_py +import subprocess + + +Path = pathlib.Path + + +class GenerateProtoGrpcCommand(setuptools.Command): + """Specialized setup command to handle kalman proto compilation. + + Generates the `kalman_pb2{_grpc}.py` files from `kalman_proto`. Assumes that + `grpc_tools.protoc` is installed. + """ + + description = "Generate `.proto` files to Python protobuf and gRPC files." + user_options = [] + + def initialize_options(self): + self.build_lib = None + + def finalize_options(self): + self.set_undefined_options("build_py", ("build_lib", "build_lib")) + + def run(self): + """Generate `kalman.proto` into `kalman_pb2{_grpc}.py`. + + This function looks more complicated than what it has to be because the + `protoc` generator is very particular in the way it generates the imports + for the generated `kalman_pb2_grpc.py` file. The final argument of the + `protoc` call has to be "mujoco_mpc/kalman.proto" in order for the import to + become `from mujoco_mpc import [kalman_pb2_proto_import]` instead of just + `import [kalman_pb2_proto_import]`. The latter would fail because the name is + meant to be relative but python3 interprets it as an absolute import. + """ + # We import here because, if the import is at the top of this file, we + # cannot resolve the dependencies without having `grpcio-tools` installed. + from grpc_tools import protoc # pylint: disable=import-outside-toplevel + + kalman_proto_filename = "kalman.proto" + kalman_proto_source_path = Path("..", "grpc", kalman_proto_filename).resolve() + assert self.build_lib is not None + build_lib_path = Path(self.build_lib).resolve() + proto_module_relative_path = Path("mujoco_mpc", "proto", kalman_proto_filename) + kalman_proto_destination_path = Path(build_lib_path, proto_module_relative_path) + kalman_proto_destination_path.parent.mkdir(parents=True, exist_ok=True) + # Copy `kalman_proto_filename` into current source. + shutil.copy(kalman_proto_source_path, kalman_proto_destination_path) + + protoc_command_parts = [ + # We use `__file__` as the first argument the same way as is done by + # `protoc` when called as `__main__` here: + # https://github.com/grpc/grpc/blob/21996c37842035661323c71b9e7040345f0915e2/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py#L172-L173. + __file__, + f"-I{build_lib_path}", + f"--python_out={build_lib_path}", + f"--grpc_python_out={build_lib_path}", + str(kalman_proto_destination_path), + ] + + protoc_returncode = protoc.main(protoc_command_parts) + + if protoc_returncode != 0: + raise subprocess.CalledProcessError( + returncode=protoc_returncode, + cmd=f"`protoc.main({protoc_command_parts})`", + ) + + self.spawn(["touch", str(kalman_proto_destination_path.parent / "__init__.py")]) + + +class CopyKalmanServerBinaryCommand(setuptools.Command): + """Specialized setup command to copy `kalman_server` next to `kalman.py`. + + Assumes that the C++ gRPC `kalman_server` binary has been manually built and + and located in the default `mujoco_mpc/build/bin` folder. + """ + + description = "Copy `kalman_server` next to `kalman.py`." + user_options = [] + + def initialize_options(self): + self.build_lib = None + + def finalize_options(self): + self.set_undefined_options("build_py", ("build_lib", "build_lib")) + + def run(self): + self._copy_binary("kalman_server") + # self._copy_binary("ui_kalman_server") + + def _copy_binary(self, binary_name): + source_path = Path(f"../build/bin/{binary_name}") + if not source_path.exists(): + raise ValueError( + f"Cannot find `{binary_name}` binary from {source_path}. Please build" + " the `{binary_name}` C++ gRPC service." + ) + assert self.build_lib is not None + build_lib_path = Path(self.build_lib).resolve() + destination_path = Path(build_lib_path, "mujoco_mpc", "mjpc", binary_name) + + self.announce(f"{source_path.resolve()=}") + self.announce(f"{destination_path.resolve()=}") + + destination_path.parent.mkdir(exist_ok=True, parents=True) + shutil.copy(source_path, destination_path) + + +class CopyTaskAssetsCommand(setuptools.Command): + """Copies `kalman_server` and `ui_kalman_server` next to `kalman.py`. + + Assumes that the C++ gRPC `kalman_server` binary has been manually built and + and located in the default `mujoco_mpc/build/bin` folder. + """ + + description = ( + "Copy task assets over to python source to make them accessible by" + " `Kalman`." + ) + user_options = [] + + def initialize_options(self): + self.build_lib = None + + def finalize_options(self): + self.set_undefined_options("build_ext", ("build_lib", "build_lib")) + + def run(self): + mjpc_tasks_path = Path(__file__).parent.parent / "mjpc" / "tasks" + source_paths = tuple(mjpc_tasks_path.rglob("*.xml")) + relative_source_paths = tuple(p.relative_to(mjpc_tasks_path) for p in source_paths) + assert self.build_lib is not None + build_lib_path = Path(self.build_lib).resolve() + destination_dir_path = Path(build_lib_path, "mujoco_mpc", "mjpc", "tasks") + self.announce( + f"Copying assets {relative_source_paths} from" + f" {mjpc_tasks_path} over to {destination_dir_path}." + ) + + for source_path, relative_source_path in zip(source_paths, relative_source_paths): + destination_path = destination_dir_path / relative_source_path + destination_path.parent.mkdir(exist_ok=True, parents=True) + shutil.copy(source_path, destination_path) + + +class BuildPyCommand(build_py.build_py): + """Specialized Python builder to handle kalman service dependencies. + + During build, this will generate the `kalman_pb2{_grpc}.py` files and copy + `kalman_server` binary next to `kalman.py`. + """ + + user_options = build_py.build_py.user_options + + def run(self): + self.run_command("generate_proto_grpc") + self.run_command("copy_task_assets") + super().run() + + +class CMakeExtension(setuptools.Extension): + """A Python extension that has been prebuilt by CMake. + + We do not want distutils to handle the build process for our extensions, so + so we pass an empty list to the super constructor. + """ + + def __init__(self, name): + super().__init__(name, sources=[]) + + +class BuildCMakeExtension(build_ext.build_ext): + """Uses CMake to build extensions.""" + + def run(self): + self._configure_and_build_kalman_server() + self.run_command("copy_kalman_server_binary") + + def _configure_and_build_kalman_server(self): + """Check for CMake.""" + cmake_command = "cmake" + build_cfg = "Debug" + mujoco_mpc_root = Path(__file__).parent.parent + mujoco_mpc_build_dir = mujoco_mpc_root / "build" + cmake_configure_args = [ + "-DCMAKE_EXPORT_COMPILE_COMMANDS:BOOL=TRUE", + f"-DCMAKE_BUILD_TYPE:STRING={build_cfg}", + "-DBUILD_TESTING:BOOL=OFF", + "-DMJPC_BUILD_GRPC_SERVICE:BOOL=ON", + ] + + if platform.system() == "Darwin" and "ARCHFLAGS" in os.environ: + osx_archs = [] + if "-arch x86_64" in os.environ["ARCHFLAGS"]: + osx_archs.append("x86_64") + if "-arch arm64" in os.environ["ARCHFLAGS"]: + osx_archs.append("arm64") + cmake_configure_args.append(f"-DCMAKE_OSX_ARCHITECTURES={';'.join(osx_archs)}") + + # TODO(hartikainen): We currently configure the builds into + # `mujoco_mpc/build`. This should use `self.build_{temp,lib}` instead, to + # isolate the Python builds from the C++ builds. + print("Configuring CMake with the following arguments:") + for arg in cmake_configure_args: + print(f" {arg}") + subprocess.check_call( + [ + cmake_command, + *cmake_configure_args, + f"-S{mujoco_mpc_root.resolve()}", + f"-B{mujoco_mpc_build_dir.resolve()}", + ], + cwd=mujoco_mpc_root, + ) + + print("Building `kalman_server` and `ui_kalman_server` with CMake") + subprocess.check_call( + [ + cmake_command, + "--build", + str(mujoco_mpc_build_dir.resolve()), + "--target", + "kalman_server", + # "ui_kalman_server", + f"-j{os.cpu_count()}", + "--config", + build_cfg, + ], + cwd=mujoco_mpc_root, + ) + + +setuptools.setup( + name="mujoco_mpc", + version="0.1.0", + author="DeepMind", + author_email="mujoco@deepmind.com", + description="MuJoCo MPC (MJPC)", + url="https://github.com/deepmind/mujoco_mpc", + license="MIT", + classifiers=[ + "Development Status :: 2 - Pre-Alpha", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: Apache Software License", + "Natural Language :: English", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Topic :: Scientific/Engineering", + ], + packages=setuptools.find_packages(), + python_requires=">=3.7", + install_requires=[ + "grpcio-tools", + "grpcio", + ], + extras_require={ + "test": [ + "absl-py", + "mujoco >= 2.3.3", + ], + }, + ext_modules=[CMakeExtension("kalman_server")], + cmdclass={ + "build_py": BuildPyCommand, + "build_ext": BuildCMakeExtension, + "generate_proto_grpc": GenerateProtoGrpcCommand, + "copy_kalman_server_binary": CopyKalmanServerBinaryCommand, + "copy_task_assets": CopyTaskAssetsCommand, + }, + package_data={ + "": [ + "mjpc/kalman_server", + # "mjpc/ui_kalman_server", + ], + }, +) \ No newline at end of file From 82b2a44376b82e52b62a7559bb32daacba52722a Mon Sep 17 00:00:00 2001 From: taylor howell Date: Tue, 15 Aug 2023 09:48:33 -0600 Subject: [PATCH 02/11] add kalman_test to build --- CMakeLists.txt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 90b72be9f..b00bf380b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -191,6 +191,6 @@ set(MJPC_COMPILE_OPTIONS "${AVX_COMPILE_OPTIONS}" "${EXTRA_COMPILE_OPTIONS}") set(MJPC_LINK_OPTIONS "${EXTRA_LINK_OPTIONS}") add_subdirectory(mjpc) -# if(MJPC_BUILD_GRPC_SERVICE) -# add_subdirectory(grpc) -# endif() +if(MJPC_BUILD_GRPC_SERVICE) + add_subdirectory(grpc) +endif() From b6b35c7c993eb6aa92aa6d49fae78f90750c571c Mon Sep 17 00:00:00 2001 From: taylor howell Date: Tue, 15 Aug 2023 09:48:56 -0600 Subject: [PATCH 03/11] add kalman_test to build --- .github/workflows/build.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e069adea0..b01e525b0 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -64,7 +64,7 @@ jobs: ${{ matrix.cmake_args }} - name: Build MuJoCo MPC working-directory: build - run: cmake --build . --config=Release ${{ matrix.cmake_build_args }} --target mjpc agent_test cost_derivatives_test norm_test rollout_test threadpool_test trajectory_test utilities_test batch_force_test batch_optimize_test batch_prior_test batch_sensor_test estimator_trajectory_test estimator_utilities_test ${{ matrix.additional_targets }} + run: cmake --build . --config=Release ${{ matrix.cmake_build_args }} --target mjpc agent_test cost_derivatives_test norm_test rollout_test threadpool_test trajectory_test utilities_test batch_force_test batch_optimize_test batch_prior_test batch_sensor_test estimator_trajectory_test estimator_utilities_test kalman_test ${{ matrix.additional_targets }} - name: Test MuJoCo MPC working-directory: build run: ctest -C Release --output-on-failure . From 459a0249a6989544cad64eef229b393181395198 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Tue, 15 Aug 2023 09:55:16 -0600 Subject: [PATCH 04/11] minor fixes --- mjpc/test/CMakeLists.txt | 16 ++++++++-------- mjpc/test/estimator/batch_prior_test.cc | 6 +++--- mjpc/test/estimator/batch_sensor_test.cc | 14 +++++++------- mjpc/test/testdata/estimator/box/task2D.xml | 8 -------- 4 files changed, 18 insertions(+), 26 deletions(-) diff --git a/mjpc/test/CMakeLists.txt b/mjpc/test/CMakeLists.txt index 8f5ba2337..112dac539 100644 --- a/mjpc/test/CMakeLists.txt +++ b/mjpc/test/CMakeLists.txt @@ -64,12 +64,12 @@ target_include_directories(simulation PRIVATE ${mujoco_SOURCE_DIR}/include gmock enable_testing() -# add_subdirectory(agent) +add_subdirectory(agent) add_subdirectory(estimator) -# add_subdirectory(gradient_planner) -# add_subdirectory(ilqg_planner) -# add_subdirectory(planners/robust) -# add_subdirectory(sampling_planner) -# add_subdirectory(state) -# add_subdirectory(tasks) -# add_subdirectory(utilities) +add_subdirectory(gradient_planner) +add_subdirectory(ilqg_planner) +add_subdirectory(planners/robust) +add_subdirectory(sampling_planner) +add_subdirectory(state) +add_subdirectory(tasks) +add_subdirectory(utilities) diff --git a/mjpc/test/estimator/batch_prior_test.cc b/mjpc/test/estimator/batch_prior_test.cc index 57fefb916..3d4d9bc9e 100644 --- a/mjpc/test/estimator/batch_prior_test.cc +++ b/mjpc/test/estimator/batch_prior_test.cc @@ -118,15 +118,15 @@ TEST(PriorCost, Particle) { // cost double cost_lambda = cost_prior(estimator.configuration.Data()); - // // gradient + // gradient FiniteDifferenceGradient fdg(nvar); fdg.Compute(cost_prior, estimator.configuration.Data(), nvar); - // // Hessian + // Hessian FiniteDifferenceHessian fdh(nvar); fdh.Compute(cost_prior, estimator.configuration.Data(), nvar); - // // ----- estimator ----- // + // ----- estimator ----- // ThreadPool pool(1); // evaluate cost diff --git a/mjpc/test/estimator/batch_sensor_test.cc b/mjpc/test/estimator/batch_sensor_test.cc index 98ec5aa67..00247a2c8 100644 --- a/mjpc/test/estimator/batch_sensor_test.cc +++ b/mjpc/test/estimator/batch_sensor_test.cc @@ -90,7 +90,7 @@ TEST(MeasurementCost, Particle) { // initialize double cost = 0.0; - // time scaling + // time scaling double time_scale = 1.0; double time_scale2 = 1.0; if (estimator.settings.time_scaling_sensor) { @@ -119,7 +119,7 @@ TEST(MeasurementCost, Particle) { // loop over sensors int shift = 0; - for (int i = 0; i < model->nsensor; i++) { + for (int i = 0; i < model->nsensor; i++) { // skip velocity, acceleration sensors (assumes position sensors are // first) if (model->sensor_needstage[i] != mjSTAGE_POS) continue; @@ -180,7 +180,7 @@ TEST(MeasurementCost, Particle) { int shift = 0; for (int i = 0; i < model->nsensor; i++) { - // sensor stage + // sensor stage int sensor_stage = model->sensor_needstage[i]; // sensor dimension @@ -189,7 +189,7 @@ TEST(MeasurementCost, Particle) { // sensor residual double* rki = rk + shift; - // time weight + // time weight double time_weight = 1.0; if (sensor_stage == mjSTAGE_VEL) { time_weight = time_scale; @@ -356,7 +356,7 @@ TEST(MeasurementCost, Box) { // initialize double cost = 0.0; - // time scale + // time scale double time_scale = 1.0; double time_scale2 = 1.0; if (estimator.settings.time_scaling_sensor) { @@ -446,7 +446,7 @@ TEST(MeasurementCost, Box) { int shift = 0; for (int i = 0; i < model->nsensor; i++) { - // sensor stage + // sensor stage int sensor_stage = model->sensor_needstage[i]; // sensor dimension @@ -455,7 +455,7 @@ TEST(MeasurementCost, Box) { // sensor residual double* rki = rk + shift; - // time weight + // time weight double time_weight = 1.0; if (sensor_stage == mjSTAGE_VEL) { time_weight = time_scale; diff --git a/mjpc/test/testdata/estimator/box/task2D.xml b/mjpc/test/testdata/estimator/box/task2D.xml index 4c69682bd..428c63623 100644 --- a/mjpc/test/testdata/estimator/box/task2D.xml +++ b/mjpc/test/testdata/estimator/box/task2D.xml @@ -46,14 +46,6 @@ - From bb9ff5ea80ba6cad3642bcd8b29c314a5d23ef40 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Thu, 17 Aug 2023 11:20:32 -0600 Subject: [PATCH 05/11] remove tblock task --- mjpc/test/testdata/estimator/tblock/task.xml | 32 -------------------- 1 file changed, 32 deletions(-) delete mode 100644 mjpc/test/testdata/estimator/tblock/task.xml diff --git a/mjpc/test/testdata/estimator/tblock/task.xml b/mjpc/test/testdata/estimator/tblock/task.xml deleted file mode 100644 index 42eefedf7..000000000 --- a/mjpc/test/testdata/estimator/tblock/task.xml +++ /dev/null @@ -1,32 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - From a8a4049d607908405d994bcdefd73c177845c98e Mon Sep 17 00:00:00 2001 From: taylor howell Date: Thu, 17 Aug 2023 11:21:51 -0600 Subject: [PATCH 06/11] remove unused batch_service_test --- grpc/batch_service_test.cc | 71 -------------------------------------- 1 file changed, 71 deletions(-) delete mode 100644 grpc/batch_service_test.cc diff --git a/grpc/batch_service_test.cc b/grpc/batch_service_test.cc deleted file mode 100644 index fc25e81d0..000000000 --- a/grpc/batch_service_test.cc +++ /dev/null @@ -1,71 +0,0 @@ -// Copyright 2023 DeepMind Technologies Limited -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Unit tests for the `BatchService` class. - -#include "grpc/batch_service.h" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "testing/base/public/gmock.h" -#include "testing/base/public/gunit.h" -#include "grpc/estimator.grpc.pb.h" -#include "grpc/estimator.pb.h" - -namespace mjpc::batch_grpc { - -using batch::grpc_gen::Batch; - -class BatchServiceTest : public ::testing::Test { - protected: - void SetUp() override { - batch_service = std::make_unique(); - grpc::ServerBuilder builder; - builder.RegisterService(batch_service.get()); - server = builder.BuildAndStart(); - std::shared_ptr channel = - server->InProcessChannel(grpc::ChannelArguments()); - stub = Batch::NewStub(channel); - } - - void TearDown() override { server->Shutdown(); } - - void RunAndCheckInit() { - grpc::ClientContext init_context; - - batch::InitRequest init_request; - - batch::InitResponse init_response; - grpc::Status init_status = - stub->Init(&init_context, init_request, &init_response); - - EXPECT_TRUE(init_status.ok()) << init_status.error_message(); - } - - std::unique_ptr batch_service; - std::unique_ptr stub; - std::unique_ptr server; -}; - -} // namespace mjpc::batch_grpc From c942d45978fca95042d59c2043449dd671cd0d7a Mon Sep 17 00:00:00 2001 From: taylor howell Date: Thu, 17 Aug 2023 11:39:52 -0600 Subject: [PATCH 07/11] minor fixes --- grpc/kalman.proto | 2 +- grpc/kalman_service.cc | 4 +-- grpc/kalman_service.h | 2 +- mjpc/estimators/batch.cc | 2 +- .../test/testdata/estimator/box/task3Drot.xml | 2 +- .../testdata/estimator/box/task3Drot2.xml | 2 +- .../testdata/estimator/particle/task_imu.xml | 2 +- python/mujoco_mpc/kalman.py | 36 ++++++++++++------- python/mujoco_mpc/kalman_test.py | 12 ++++--- python/setup_batch.py | 4 +-- python/setup_kalman.py | 31 +++++++++++----- 11 files changed, 63 insertions(+), 36 deletions(-) diff --git a/grpc/kalman.proto b/grpc/kalman.proto index 3e6e24bb2..a7a62970f 100644 --- a/grpc/kalman.proto +++ b/grpc/kalman.proto @@ -119,4 +119,4 @@ message NoiseRequest { message NoiseResponse { Noise noise = 2; -} \ No newline at end of file +} diff --git a/grpc/kalman_service.cc b/grpc/kalman_service.cc index 5291b15f1..58feacd4f 100644 --- a/grpc/kalman_service.cc +++ b/grpc/kalman_service.cc @@ -114,8 +114,8 @@ grpc::Status KalmanService::Init(grpc::ServerContext* context, } // move - kalman_model_override_ = std::move(tmp_model); - mjModel* model = kalman_model_override_.get(); + model_override_ = std::move(tmp_model); + mjModel* model = model_override_.get(); // initialize kalman kalman_.Initialize(model); diff --git a/grpc/kalman_service.h b/grpc/kalman_service.h index 63d3ed78d..fca9e2f57 100644 --- a/grpc/kalman_service.h +++ b/grpc/kalman_service.h @@ -80,7 +80,7 @@ class KalmanService final : public kalman::Kalman::Service { // kalman mjpc::Kalman kalman_; - mjpc::UniqueMjModel kalman_model_override_ = {nullptr, mj_deleteModel}; + mjpc::UniqueMjModel model_override_ = {nullptr, mj_deleteModel}; // threadpool mjpc::ThreadPool thread_pool_; diff --git a/mjpc/estimators/batch.cc b/mjpc/estimators/batch.cc index 657256ae1..d83e43715 100644 --- a/mjpc/estimators/batch.cc +++ b/mjpc/estimators/batch.cc @@ -2764,7 +2764,7 @@ std::string StatusString(int code) { case kSmallDirectionFailure: return "SMALL_DIRECTION_FAILURE"; case kMaxRegularizationFailure: - return "kMaxBatchRegularization_FAILURE"; + return "MAX_REGULARIZATION_FAILURE"; case kCostDifferenceFailure: return "COST_DIFFERENCE_FAILURE"; case kExpectedDecreaseFailure: diff --git a/mjpc/test/testdata/estimator/box/task3Drot.xml b/mjpc/test/testdata/estimator/box/task3Drot.xml index aab3b9ec3..448c1f1fe 100644 --- a/mjpc/test/testdata/estimator/box/task3Drot.xml +++ b/mjpc/test/testdata/estimator/box/task3Drot.xml @@ -62,4 +62,4 @@ - \ No newline at end of file + diff --git a/mjpc/test/testdata/estimator/box/task3Drot2.xml b/mjpc/test/testdata/estimator/box/task3Drot2.xml index 59d2cdd70..8a14a7ca3 100644 --- a/mjpc/test/testdata/estimator/box/task3Drot2.xml +++ b/mjpc/test/testdata/estimator/box/task3Drot2.xml @@ -72,4 +72,4 @@ - \ No newline at end of file + diff --git a/mjpc/test/testdata/estimator/particle/task_imu.xml b/mjpc/test/testdata/estimator/particle/task_imu.xml index d0835cc64..611f63b38 100644 --- a/mjpc/test/testdata/estimator/particle/task_imu.xml +++ b/mjpc/test/testdata/estimator/particle/task_imu.xml @@ -53,4 +53,4 @@ - \ No newline at end of file + diff --git a/python/mujoco_mpc/kalman.py b/python/mujoco_mpc/kalman.py index 528b5babc..834243a45 100644 --- a/python/mujoco_mpc/kalman.py +++ b/python/mujoco_mpc/kalman.py @@ -76,10 +76,12 @@ def __init__( [str(server_binary_path), f"--mjpc_port={self.port}"], stdout=subprocess.PIPE if colab_logging else None, ) - os.set_blocking(self.server_process.stdout.fileno(), False) + # os.set_blocking(self.server_process.stdout.fileno(), False) atexit.register(self.server_process.kill) - credentials = grpc.local_channel_credentials(grpc.LocalConnectionType.LOCAL_TCP) + credentials = grpc.local_channel_credentials( + grpc.LocalConnectionType.LOCAL_TCP + ) self.channel = grpc.secure_channel(f"localhost:{self.port}", credentials) grpc.channel_ready_future(self.channel).result(timeout=10) self.stub = kalman_pb2_grpc.KalmanStub(self.channel) @@ -172,7 +174,9 @@ def settings( } def update_measurement( - self, ctrl: Optional[npt.ArrayLike] = [], sensor: Optional[npt.ArrayLike] = [] + self, + ctrl: Optional[npt.ArrayLike] = [], + sensor: Optional[npt.ArrayLike] = [], ): # request request = kalman_pb2.UpdateMeasurementRequest( @@ -218,7 +222,9 @@ def state(self, state: Optional[npt.ArrayLike] = []) -> np.ndarray: # return state return np.array(response.state) - def covariance(self, covariance: Optional[npt.ArrayLike] = None) -> np.ndarray: + def covariance( + self, covariance: Optional[npt.ArrayLike] = None + ) -> np.ndarray: # input inputs = kalman_pb2.Covariance( covariance=covariance.flatten() if covariance is not None else None, @@ -233,10 +239,14 @@ def covariance(self, covariance: Optional[npt.ArrayLike] = None) -> np.ndarray: response = self._wait(self.stub.Covariance.future(request)).covariance # return covariance - return np.array(response.covariance).reshape(response.dimension, response.dimension) + return np.array(response.covariance).reshape( + response.dimension, response.dimension + ) def noise( - self, process: Optional[npt.ArrayLike] = [], sensor: Optional[npt.ArrayLike] = [] + self, + process: Optional[npt.ArrayLike] = [], + sensor: Optional[npt.ArrayLike] = [], ) -> dict[str, np.ndarray]: # inputs inputs = kalman_pb2.Noise( @@ -260,11 +270,11 @@ def noise( def _wait(self, future): """Waits for the future to complete, while printing out subprocess stdout.""" - if self._colab_logging: - while True: - line = self.server_process.stdout.readline() - if line: - sys.stdout.write(line.decode("utf-8")) - if future.done(): - break + # if self._colab_logging: + # while True: + # line = self.server_process.stdout.readline() + # if line: + # sys.stdout.write(line.decode("utf-8")) + # if future.done(): + # break return future.result() diff --git a/python/mujoco_mpc/kalman_test.py b/python/mujoco_mpc/kalman_test.py index dcab43586..a95b5ef2e 100644 --- a/python/mujoco_mpc/kalman_test.py +++ b/python/mujoco_mpc/kalman_test.py @@ -63,12 +63,16 @@ def test_updates(self): # covariance nvelocity = 2 * model.nv - F = np.random.normal(scale=1.0, size=(nvelocity**2)).reshape(nvelocity, nvelocity) + F = np.random.normal(scale=1.0, size=(nvelocity**2)).reshape( + nvelocity, nvelocity + ) covariance = F.T @ F covariance_response = kalman.covariance(covariance=covariance) # test covariance - self.assertLess(np.linalg.norm((covariance_response - covariance).ravel()), 1.0e-5) + self.assertLess( + np.linalg.norm((covariance_response - covariance).ravel()), 1.0e-5 + ) self.assertTrue(covariance_response.shape == (nvelocity, nvelocity)) # noise @@ -91,8 +95,8 @@ def test_updates(self): # timers timer = kalman.timers() - self.assertTrue(timer["measurement"] > 0.0) - self.assertTrue(timer["prediction"] > 0.0) + self.assertLess(0.0, timer["measurement"]) + self.assertLess(0.0, timer["prediction"]) if __name__ == "__main__": diff --git a/python/setup_batch.py b/python/setup_batch.py index cbc45767a..bd07a3b66 100644 --- a/python/setup_batch.py +++ b/python/setup_batch.py @@ -59,9 +59,7 @@ def run(self): from grpc_tools import protoc # pylint: disable=import-outside-toplevel batch_proto_filename = "batch.proto" - batch_proto_source_path = Path( - "..", "grpc", batch_proto_filename - ).resolve() + batch_proto_source_path = Path("..", "grpc", batch_proto_filename).resolve() assert self.build_lib is not None build_lib_path = Path(self.build_lib).resolve() proto_module_relative_path = Path( diff --git a/python/setup_kalman.py b/python/setup_kalman.py index 258d4c48e..7e1b5c9f4 100644 --- a/python/setup_kalman.py +++ b/python/setup_kalman.py @@ -59,11 +59,17 @@ def run(self): from grpc_tools import protoc # pylint: disable=import-outside-toplevel kalman_proto_filename = "kalman.proto" - kalman_proto_source_path = Path("..", "grpc", kalman_proto_filename).resolve() + kalman_proto_source_path = Path( + "..", "grpc", kalman_proto_filename + ).resolve() assert self.build_lib is not None build_lib_path = Path(self.build_lib).resolve() - proto_module_relative_path = Path("mujoco_mpc", "proto", kalman_proto_filename) - kalman_proto_destination_path = Path(build_lib_path, proto_module_relative_path) + proto_module_relative_path = Path( + "mujoco_mpc", "proto", kalman_proto_filename + ) + kalman_proto_destination_path = Path( + build_lib_path, proto_module_relative_path + ) kalman_proto_destination_path.parent.mkdir(parents=True, exist_ok=True) # Copy `kalman_proto_filename` into current source. shutil.copy(kalman_proto_source_path, kalman_proto_destination_path) @@ -87,7 +93,9 @@ def run(self): cmd=f"`protoc.main({protoc_command_parts})`", ) - self.spawn(["touch", str(kalman_proto_destination_path.parent / "__init__.py")]) + self.spawn( + ["touch", str(kalman_proto_destination_path.parent / "__init__.py")] + ) class CopyKalmanServerBinaryCommand(setuptools.Command): @@ -136,7 +144,8 @@ class CopyTaskAssetsCommand(setuptools.Command): """ description = ( - "Copy task assets over to python source to make them accessible by" " `Kalman`." + "Copy task assets over to python source to make them accessible by" + " `Kalman`." ) user_options = [] @@ -149,7 +158,9 @@ def finalize_options(self): def run(self): mjpc_tasks_path = Path(__file__).parent.parent / "mjpc" / "tasks" source_paths = tuple(mjpc_tasks_path.rglob("*.xml")) - relative_source_paths = tuple(p.relative_to(mjpc_tasks_path) for p in source_paths) + relative_source_paths = tuple( + p.relative_to(mjpc_tasks_path) for p in source_paths + ) assert self.build_lib is not None build_lib_path = Path(self.build_lib).resolve() destination_dir_path = Path(build_lib_path, "mujoco_mpc", "mjpc", "tasks") @@ -158,7 +169,9 @@ def run(self): f" {mjpc_tasks_path} over to {destination_dir_path}." ) - for source_path, relative_source_path in zip(source_paths, relative_source_paths): + for source_path, relative_source_path in zip( + source_paths, relative_source_paths + ): destination_path = destination_dir_path / relative_source_path destination_path.parent.mkdir(exist_ok=True, parents=True) shutil.copy(source_path, destination_path) @@ -216,7 +229,9 @@ def _configure_and_build_kalman_server(self): osx_archs.append("x86_64") if "-arch arm64" in os.environ["ARCHFLAGS"]: osx_archs.append("arm64") - cmake_configure_args.append(f"-DCMAKE_OSX_ARCHITECTURES={';'.join(osx_archs)}") + cmake_configure_args.append( + f"-DCMAKE_OSX_ARCHITECTURES={';'.join(osx_archs)}" + ) # TODO(hartikainen): We currently configure the builds into # `mujoco_mpc/build`. This should use `self.build_{temp,lib}` instead, to From a12be8051eb019cb7adf5d1f23dcc940513c46f0 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Thu, 17 Aug 2023 11:43:37 -0600 Subject: [PATCH 08/11] reformat setup_kalman.py --- python/setup_kalman.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/setup_kalman.py b/python/setup_kalman.py index 7e1b5c9f4..14a42e979 100644 --- a/python/setup_kalman.py +++ b/python/setup_kalman.py @@ -23,7 +23,6 @@ from setuptools.command import build_py import subprocess - Path = pathlib.Path From 5b9038e538167feaa5769b5c6fce335844c3290a Mon Sep 17 00:00:00 2001 From: taylor howell Date: Thu, 17 Aug 2023 11:51:55 -0600 Subject: [PATCH 09/11] more minor fixes --- grpc/CMakeLists.txt | 2 +- grpc/kalman_server.cc | 1 - mjpc/estimators/estimator.h | 4 ++-- mjpc/estimators/gui.cc | 4 ++-- mjpc/estimators/gui.h | 4 ++-- mjpc/estimators/include.cc | 2 ++ 6 files changed, 9 insertions(+), 8 deletions(-) diff --git a/grpc/CMakeLists.txt b/grpc/CMakeLists.txt index c0cbf465d..8c9e79683 100644 --- a/grpc/CMakeLists.txt +++ b/grpc/CMakeLists.txt @@ -264,4 +264,4 @@ target_link_libraries( target_include_directories(kalman_server PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..) message(KALMAN_SERVICE_COMPILE_OPTIONS=${KALMAN_SERVICE_COMPILE_OPTIONS}) target_compile_options(kalman_server PUBLIC ${KALMAN_SERVICE_COMPILE_OPTIONS}) -target_link_options(kalman_server PRIVATE ${KALMAN_SERVICE_LINK_OPTIONS}) \ No newline at end of file +target_link_options(kalman_server PRIVATE ${KALMAN_SERVICE_LINK_OPTIONS}) diff --git a/grpc/kalman_server.cc b/grpc/kalman_server.cc index 363deba57..a6900a448 100644 --- a/grpc/kalman_server.cc +++ b/grpc/kalman_server.cc @@ -22,7 +22,6 @@ #include #include #include - // DEEPMIND INTERNAL IMPORT #include #include diff --git a/mjpc/estimators/estimator.h b/mjpc/estimators/estimator.h index bf9552f60..4a8c780aa 100644 --- a/mjpc/estimators/estimator.h +++ b/mjpc/estimators/estimator.h @@ -15,12 +15,12 @@ #ifndef MJPC_ESTIMATORS_ESTIMATOR_H_ #define MJPC_ESTIMATORS_ESTIMATOR_H_ -#include - #include #include #include +#include + #include "mjpc/estimators/gui.h" #include "mjpc/utilities.h" diff --git a/mjpc/estimators/gui.cc b/mjpc/estimators/gui.cc index 53be837b6..9e1b4dc9b 100644 --- a/mjpc/estimators/gui.cc +++ b/mjpc/estimators/gui.cc @@ -14,10 +14,10 @@ #include "mjpc/estimators/gui.h" -#include - #include +#include + #include "mjpc/utilities.h" namespace mjpc { diff --git a/mjpc/estimators/gui.h b/mjpc/estimators/gui.h index b5f6b36b7..e773e1e87 100644 --- a/mjpc/estimators/gui.h +++ b/mjpc/estimators/gui.h @@ -15,10 +15,10 @@ #ifndef MJPC_ESTIMATORS_GUI_H_ #define MJPC_ESTIMATORS_GUI_H_ -#include - #include +#include + namespace mjpc { // data that is modified by the GUI and later set in Estimator diff --git a/mjpc/estimators/include.cc b/mjpc/estimators/include.cc index 4a0641e8e..6a824fe8f 100644 --- a/mjpc/estimators/include.cc +++ b/mjpc/estimators/include.cc @@ -17,6 +17,7 @@ namespace mjpc { const char kEstimatorNames[] = + "Ground Truth\n" "Kalman\n" "Batch"; @@ -28,6 +29,7 @@ std::vector> LoadEstimators() { // add estimators estimators.emplace_back(new mjpc::GroundTruth()); // ground truth state estimators.emplace_back(new mjpc::Kalman()); // extended Kalman filter + estimators.emplace_back(new mjpc::Batch(1)); // recursive batch filter return estimators; } From 704ba9e1750b3b55ce7c5d6caa0bd78386228a68 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Thu, 17 Aug 2023 11:52:57 -0600 Subject: [PATCH 10/11] reformat kalman.h --- mjpc/estimators/kalman.h | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/mjpc/estimators/kalman.h b/mjpc/estimators/kalman.h index 405e6ff0b..baf23f946 100644 --- a/mjpc/estimators/kalman.h +++ b/mjpc/estimators/kalman.h @@ -99,9 +99,7 @@ class Kalman : public Estimator { }; // set time - void SetTime(double time) override { - this->time = time; - } + void SetTime(double time) override { this->time = time; } // set covariance void SetCovariance(const double* covariance) override { From 3558421ef023a3abd940db47dc508eedcd9b0295 Mon Sep 17 00:00:00 2001 From: taylor howell Date: Sun, 20 Aug 2023 11:32:47 -0600 Subject: [PATCH 11/11] remove kalman grpc + python --- grpc/CMakeLists.txt | 57 ------ grpc/kalman.proto | 122 ------------ grpc/kalman_server.cc | 57 ------ grpc/kalman_service.cc | 314 ------------------------------- grpc/kalman_service.h | 91 --------- python/mujoco_mpc/kalman.py | 280 --------------------------- python/mujoco_mpc/kalman_test.py | 103 ---------- python/setup_kalman.py | 312 ------------------------------ 8 files changed, 1336 deletions(-) delete mode 100644 grpc/kalman.proto delete mode 100644 grpc/kalman_server.cc delete mode 100644 grpc/kalman_service.cc delete mode 100644 grpc/kalman_service.h delete mode 100644 python/mujoco_mpc/kalman.py delete mode 100644 python/mujoco_mpc/kalman_test.py delete mode 100644 python/setup_kalman.py diff --git a/grpc/CMakeLists.txt b/grpc/CMakeLists.txt index 8c9e79683..d10df47f2 100644 --- a/grpc/CMakeLists.txt +++ b/grpc/CMakeLists.txt @@ -55,8 +55,6 @@ get_filename_component(agent_service_proto "./agent.proto" ABSOLUTE) get_filename_component(agent_service_proto_path "${agent_service_proto}" PATH) get_filename_component(batch_service_proto "./batch.proto" ABSOLUTE) get_filename_component(batch_service_proto_path "${batch_service_proto}" PATH) -get_filename_component(kalman_service_proto "./kalman.proto" ABSOLUTE) -get_filename_component(kalman_service_proto_path "${kalman_service_proto}" PATH) # Generated sources set(agent_service_proto_srcs "${CMAKE_CURRENT_BINARY_DIR}/agent.pb.cc") @@ -69,11 +67,6 @@ set(batch_service_proto_hdrs "${CMAKE_CURRENT_BINARY_DIR}/batch.pb.h") set(batch_service_grpc_srcs "${CMAKE_CURRENT_BINARY_DIR}/batch.grpc.pb.cc") set(batch_service_grpc_hdrs "${CMAKE_CURRENT_BINARY_DIR}/batch.grpc.pb.h") -set(kalman_service_proto_srcs "${CMAKE_CURRENT_BINARY_DIR}/kalman.pb.cc") -set(kalman_service_proto_hdrs "${CMAKE_CURRENT_BINARY_DIR}/kalman.pb.h") -set(kalman_service_grpc_srcs "${CMAKE_CURRENT_BINARY_DIR}/kalman.grpc.pb.cc") -set(kalman_service_grpc_hdrs "${CMAKE_CURRENT_BINARY_DIR}/kalman.grpc.pb.h") - message("We need the following for agent/batch_protos:") message(_GRPC_CPP_PLUGIN_EXECUTABLE${_GRPC_CPP_PLUGIN_EXECUTABLE}) @@ -111,23 +104,6 @@ add_custom_command( "${batch_service_proto}" ) -add_custom_command( - OUTPUT - "${kalman_service_proto_srcs}" - "${kalman_service_proto_hdrs}" - "${kalman_service_grpc_srcs}" - "${kalman_service_grpc_hdrs}" - COMMAND ${_PROTOBUF_PROTOC} - ARGS - --grpc_out "${CMAKE_CURRENT_BINARY_DIR}" - --cpp_out "${CMAKE_CURRENT_BINARY_DIR}" - -I "${kalman_service_proto_path}" - --plugin=protoc-gen-grpc="${_GRPC_CPP_PLUGIN_EXECUTABLE}" - "${kalman_service_proto}" - DEPENDS - "${kalman_service_proto}" -) - # Include generated *.pb.h files include_directories(${CMAKE_CURRENT_BINARY_DIR}/..) @@ -232,36 +208,3 @@ target_include_directories(batch_server PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..) message(BATCH_SERVICE_COMPILE_OPTIONS=${BATCH_SERVICE_COMPILE_OPTIONS}) target_compile_options(batch_server PUBLIC ${BATCH_SERVICE_COMPILE_OPTIONS}) target_link_options(batch_server PRIVATE ${BATCH_SERVICE_LINK_OPTIONS}) - -add_executable( - kalman_server - kalman_server.cc - ${kalman_service_grpc_srcs} - ${kalman_service_grpc_hdrs} - ${kalman_service_proto_srcs} - ${kalman_service_proto_hdrs} - kalman_service.h - kalman_service.cc -) - -target_link_libraries( - kalman_server - # kalman_service_grpc_proto - absl::check - absl::flags - absl::flags_parse - absl::log - absl::random_random - absl::status - absl::strings - ${_REFLECTION} - ${_GRPC_GRPCPP} - ${_PROTOBUF_LIBPROTOBUF} - mujoco::mujoco - libmjpc -) - -target_include_directories(kalman_server PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..) -message(KALMAN_SERVICE_COMPILE_OPTIONS=${KALMAN_SERVICE_COMPILE_OPTIONS}) -target_compile_options(kalman_server PUBLIC ${KALMAN_SERVICE_COMPILE_OPTIONS}) -target_link_options(kalman_server PRIVATE ${KALMAN_SERVICE_LINK_OPTIONS}) diff --git a/grpc/kalman.proto b/grpc/kalman.proto deleted file mode 100644 index a7a62970f..000000000 --- a/grpc/kalman.proto +++ /dev/null @@ -1,122 +0,0 @@ -// Copyright 2023 DeepMind Technologies Limited -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -syntax = "proto3"; - -package kalman; - -service Kalman { - // Initialize Kalman - rpc Init(InitRequest) returns (InitResponse); - // Reset Kalman - rpc Reset(ResetRequest) returns (ResetResponse); - // Kalman settings - rpc Settings(SettingsRequest) returns (SettingsResponse); - // Kalman measurement update - rpc UpdateMeasurement(UpdateMeasurementRequest) returns (UpdateMeasurementResponse); - // Kalman prediction update - rpc UpdatePrediction(UpdatePredictionRequest) returns (UpdatePredictionResponse); - // Kalman timers - rpc Timers(TimersRequest) returns (TimersResponse); - // Kalman state - rpc State(StateRequest) returns (StateResponse); - // Kalman covariance - rpc Covariance(CovarianceRequest) returns (CovarianceResponse); - // Kalman noise - rpc Noise(NoiseRequest) returns (NoiseResponse); -} - -message MjModel { - optional bytes mjb = 1; - optional string xml = 2; -} - -message InitRequest { - optional MjModel model = 1; -} - -message InitResponse {} - -message Settings { - optional double epsilon = 1; - optional bool flg_centered = 2; -} - -message SettingsRequest { - optional Settings settings = 1; -} - -message SettingsResponse { - Settings settings = 1; -} - -message ResetRequest {} - -message ResetResponse {} - -message UpdateMeasurementRequest { - repeated double ctrl = 1 [packed = true]; - repeated double sensor = 2 [packed = true]; -} - -message UpdateMeasurementResponse {} - -message UpdatePredictionRequest {} - -message UpdatePredictionResponse {} - -message TimersRequest {} - -message TimersResponse { - double measurement = 1; - double prediction = 2; -} - -message State { - repeated double state = 1 [packed = true]; -} - -message StateRequest { - optional State state = 1; -} - -message StateResponse { - State state = 1; -} - -message Covariance { - repeated double covariance = 1 [packed = true]; - optional int32 dimension = 2; -} - -message CovarianceRequest { - optional Covariance covariance = 1; -} - -message CovarianceResponse { - Covariance covariance = 1; -} - -message Noise { - repeated double process = 1 [packed = true]; - repeated double sensor = 2 [packed = true]; -} - -message NoiseRequest { - optional Noise noise = 1; -} - -message NoiseResponse { - Noise noise = 2; -} diff --git a/grpc/kalman_server.cc b/grpc/kalman_server.cc deleted file mode 100644 index a6900a448..000000000 --- a/grpc/kalman_server.cc +++ /dev/null @@ -1,57 +0,0 @@ -// Copyright 2023 DeepMind Technologies Limited -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// Startup code for `Kalman` server. - -#include -#include -#include - -#include -#include -#include -#include -// DEEPMIND INTERNAL IMPORT -#include -#include -#include - -#include "grpc/kalman_service.h" - -ABSL_FLAG(int32_t, mjpc_port, 10000, "port to listen on"); - -int main(int argc, char** argv) { - absl::ParseCommandLine(argc, argv); - absl::ParseCommandLine(argc, argv); - int port = absl::GetFlag(FLAGS_mjpc_port); - - std::string server_address = absl::StrCat("[::]:", port); - - std::shared_ptr server_credentials = - grpc::experimental::LocalServerCredentials(LOCAL_TCP); - grpc::ServerBuilder builder; - builder.AddListeningPort(server_address, server_credentials); - - kalman_grpc::KalmanService service; - builder.SetMaxReceiveMessageSize(40 * 1024 * 1024); - builder.RegisterService(&service); - - std::unique_ptr server(builder.BuildAndStart()); - LOG(INFO) << "Server listening on " << server_address; - - // Keep the program running until the server shuts down. - server->Wait(); - - return 0; -} diff --git a/grpc/kalman_service.cc b/grpc/kalman_service.cc deleted file mode 100644 index 58feacd4f..000000000 --- a/grpc/kalman_service.cc +++ /dev/null @@ -1,314 +0,0 @@ -// Copyright 2023 DeepMind Technologies Limited -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include "grpc/kalman_service.h" - -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "grpc/kalman.pb.h" -#include "mjpc/estimators/kalman.h" - -namespace kalman_grpc { - -using ::kalman::CovarianceRequest; -using ::kalman::CovarianceResponse; -using ::kalman::InitRequest; -using ::kalman::InitResponse; -using ::kalman::NoiseRequest; -using ::kalman::NoiseResponse; -using ::kalman::ResetRequest; -using ::kalman::ResetResponse; -using ::kalman::SettingsRequest; -using ::kalman::SettingsResponse; -using ::kalman::StateRequest; -using ::kalman::StateResponse; -using ::kalman::TimersRequest; -using ::kalman::TimersResponse; -using ::kalman::UpdateMeasurementRequest; -using ::kalman::UpdateMeasurementResponse; -using ::kalman::UpdatePredictionRequest; -using ::kalman::UpdatePredictionResponse; - -// TODO(taylor): make CheckSize utility function -namespace { -absl::Status CheckSize(std::string_view name, int model_size, int vector_size) { - std::ostringstream error_string; - if (model_size != vector_size) { - error_string << "expected " << name << " size " << model_size << ", got " - << vector_size; - return absl::InvalidArgumentError(error_string.str()); - } - return absl::OkStatus(); -} -} // namespace - -#define CHECK_SIZE(name, n1, n2) \ - { \ - auto expr = (CheckSize(name, n1, n2)); \ - if (!(expr).ok()) { \ - return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, \ - (expr).ToString()); \ - } \ - } - -KalmanService::~KalmanService() {} - -grpc::Status KalmanService::Init(grpc::ServerContext* context, - const kalman::InitRequest* request, - kalman::InitResponse* response) { - // ----- initialize with model ----- // - mjpc::UniqueMjModel tmp_model = {nullptr, mj_deleteModel}; - - // convert message - if (request->has_model() && request->model().has_mjb()) { - std::string_view mjb = request->model().mjb(); - static constexpr char file[] = "temporary-filename.mjb"; - // mjVFS structs need to be allocated on the heap, because it's ~2MB - auto vfs = std::make_unique(); - mj_defaultVFS(vfs.get()); - mj_makeEmptyFileVFS(vfs.get(), file, mjb.size()); - int file_idx = mj_findFileVFS(vfs.get(), file); - memcpy(vfs->filedata[file_idx], mjb.data(), mjb.size()); - tmp_model = {mj_loadModel(file, vfs.get()), mj_deleteModel}; - mj_deleteFileVFS(vfs.get(), file); - } else if (request->has_model() && request->model().has_xml()) { - std::string_view model_xml = request->model().xml(); - char load_error[1024] = ""; - - // TODO(taylor): utilize grpc_agent_util method - static constexpr char file[] = "temporary-filename.xml"; - // mjVFS structs need to be allocated on the heap, because it's ~2MB - auto vfs = std::make_unique(); - mj_defaultVFS(vfs.get()); - mj_makeEmptyFileVFS(vfs.get(), file, model_xml.size()); - int file_idx = mj_findFileVFS(vfs.get(), file); - memcpy(vfs->filedata[file_idx], model_xml.data(), model_xml.size()); - tmp_model = {mj_loadXML(file, vfs.get(), load_error, sizeof(load_error)), - mj_deleteModel}; - mj_deleteFileVFS(vfs.get(), file); - } else { - mju_error("Failed to create mjModel."); - } - - // move - model_override_ = std::move(tmp_model); - mjModel* model = model_override_.get(); - - // initialize kalman - kalman_.Initialize(model); - kalman_.Reset(); - - return grpc::Status::OK; -} - -grpc::Status KalmanService::Reset(grpc::ServerContext* context, - const kalman::ResetRequest* request, - kalman::ResetResponse* response) { - if (!Initialized()) { - return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; - } - - // reset - kalman_.Reset(); - - return grpc::Status::OK; -} - -grpc::Status KalmanService::Settings(grpc::ServerContext* context, - const kalman::SettingsRequest* request, - kalman::SettingsResponse* response) { - if (!Initialized()) { - return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; - } - - // settings - kalman::Settings input = request->settings(); - kalman::Settings* output = response->mutable_settings(); - - // epsilon - if (input.has_epsilon()) { - kalman_.settings.epsilon = input.epsilon(); - } - output->set_epsilon(kalman_.settings.epsilon); - - // flg_centered - if (input.has_flg_centered()) { - kalman_.settings.flg_centered = input.flg_centered(); - } - output->set_flg_centered(kalman_.settings.flg_centered); - - return grpc::Status::OK; -} - -grpc::Status KalmanService::UpdateMeasurement( - grpc::ServerContext* context, - const kalman::UpdateMeasurementRequest* request, - kalman::UpdateMeasurementResponse* response) { - if (!Initialized()) { - return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; - } - - // measurement update - kalman_.UpdateMeasurement(request->ctrl().data(), request->sensor().data()); - - return grpc::Status::OK; -} - -grpc::Status KalmanService::UpdatePrediction( - grpc::ServerContext* context, - const kalman::UpdatePredictionRequest* request, - kalman::UpdatePredictionResponse* response) { - if (!Initialized()) { - return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; - } - - // prediction update - kalman_.UpdatePrediction(); - - return grpc::Status::OK; -} - -grpc::Status KalmanService::Timers(grpc::ServerContext* context, - const kalman::TimersRequest* request, - kalman::TimersResponse* response) { - if (!Initialized()) { - return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; - } - - // measurement - response->set_measurement(kalman_.TimerMeasurement()); - - // prediction - response->set_prediction(kalman_.TimerPrediction()); - - return grpc::Status::OK; -} - -grpc::Status KalmanService::State(grpc::ServerContext* context, - const kalman::StateRequest* request, - kalman::StateResponse* response) { - if (!Initialized()) { - return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; - } - - // unpack input/output - kalman::State input = request->state(); - kalman::State* output = response->mutable_state(); - - // set state - int nstate = kalman_.model->nq + kalman_.model->nv; - if (input.state_size() > 0) { - CHECK_SIZE("state", nstate, input.state_size()); - mju_copy(kalman_.state.data(), input.state().data(), nstate); - } - - // get state - double* state = kalman_.state.data(); - for (int i = 0; i < nstate; i++) { - output->add_state(state[i]); - } - - return grpc::Status::OK; -} - -grpc::Status KalmanService::Covariance(grpc::ServerContext* context, - const kalman::CovarianceRequest* request, - kalman::CovarianceResponse* response) { - if (!Initialized()) { - return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; - } - - // unpack input/output - kalman::Covariance input = request->covariance(); - kalman::Covariance* output = response->mutable_covariance(); - - // dimensions - int nvelocity = 2 * kalman_.model->nv; - int ncovariance = nvelocity * nvelocity; - - // set dimension - output->set_dimension(nvelocity); - - // set covariance - if (input.covariance_size() > 0) { - CHECK_SIZE("covariance", ncovariance, input.covariance_size()); - mju_copy(kalman_.covariance.data(), input.covariance().data(), ncovariance); - } - - // get covariance - double* covariance = kalman_.covariance.data(); - for (int i = 0; i < ncovariance; i++) { - output->add_covariance(covariance[i]); - } - - return grpc::Status::OK; -} - -grpc::Status KalmanService::Noise(grpc::ServerContext* context, - const kalman::NoiseRequest* request, - kalman::NoiseResponse* response) { - if (!Initialized()) { - return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."}; - } - - // unpack input/output - kalman::Noise input = request->noise(); - kalman::Noise* output = response->mutable_noise(); - - // dimensions - int nprocess = 2 * kalman_.model->nv; - int nsensor = kalman_.model->nsensordata; - - // set process noise - if (input.process_size() > 0) { - CHECK_SIZE("process noise", nprocess, input.process_size()); - mju_copy(kalman_.noise_process.data(), input.process().data(), nprocess); - } - - // get process noise - double* process = kalman_.noise_process.data(); - for (int i = 0; i < nprocess; i++) { - output->add_process(process[i]); - } - - // set sensor noise - if (input.sensor_size() > 0) { - CHECK_SIZE("sensor noise", nsensor, input.sensor_size()); - mju_copy(kalman_.noise_sensor.data(), input.sensor().data(), nsensor); - } - - // get sensor noise - double* sensor = kalman_.noise_sensor.data(); - for (int i = 0; i < nsensor; i++) { - output->add_sensor(sensor[i]); - } - - return grpc::Status::OK; -} - -#undef CHECK_SIZE - -} // namespace kalman_grpc diff --git a/grpc/kalman_service.h b/grpc/kalman_service.h deleted file mode 100644 index fca9e2f57..000000000 --- a/grpc/kalman_service.h +++ /dev/null @@ -1,91 +0,0 @@ -// Copyright 2023 DeepMind Technologies Limited -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// An implementation of the `Kalman` gRPC service. - -#ifndef GRPC_KALMAN_SERVICE_H_ -#define GRPC_KALMAN_SERVICE_H_ - -#include -#include - -#include -#include -#include - -#include "grpc/kalman.grpc.pb.h" -#include "grpc/kalman.pb.h" -#include "mjpc/estimators/kalman.h" -#include "mjpc/threadpool.h" -#include "mjpc/utilities.h" - -namespace kalman_grpc { - -class KalmanService final : public kalman::Kalman::Service { - public: - explicit KalmanService() : thread_pool_(1) {} - ~KalmanService(); - - grpc::Status Init(grpc::ServerContext* context, - const kalman::InitRequest* request, - kalman::InitResponse* response) override; - - grpc::Status Reset(grpc::ServerContext* context, - const kalman::ResetRequest* request, - kalman::ResetResponse* response) override; - - grpc::Status Settings(grpc::ServerContext* context, - const kalman::SettingsRequest* request, - kalman::SettingsResponse* response) override; - - grpc::Status UpdateMeasurement( - grpc::ServerContext* context, - const kalman::UpdateMeasurementRequest* request, - kalman::UpdateMeasurementResponse* response) override; - - grpc::Status UpdatePrediction( - grpc::ServerContext* context, - const kalman::UpdatePredictionRequest* request, - kalman::UpdatePredictionResponse* response) override; - - grpc::Status Timers(grpc::ServerContext* context, - const kalman::TimersRequest* request, - kalman::TimersResponse* response) override; - - grpc::Status State(grpc::ServerContext* context, - const kalman::StateRequest* request, - kalman::StateResponse* response) override; - - grpc::Status Covariance(grpc::ServerContext* context, - const kalman::CovarianceRequest* request, - kalman::CovarianceResponse* response) override; - - grpc::Status Noise(grpc::ServerContext* context, - const kalman::NoiseRequest* request, - kalman::NoiseResponse* response) override; - - private: - bool Initialized() const { return kalman_.model; } - - // kalman - mjpc::Kalman kalman_; - mjpc::UniqueMjModel model_override_ = {nullptr, mj_deleteModel}; - - // threadpool - mjpc::ThreadPool thread_pool_; -}; - -} // namespace kalman_grpc - -#endif // GRPC_KALMAN_SERVICE_H_ diff --git a/python/mujoco_mpc/kalman.py b/python/mujoco_mpc/kalman.py deleted file mode 100644 index 834243a45..000000000 --- a/python/mujoco_mpc/kalman.py +++ /dev/null @@ -1,280 +0,0 @@ -# Copyright 2023 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -"""Python interface for interface with Kalman.""" - -import atexit -import os -import pathlib -import socket -import subprocess -import sys -import tempfile -from typing import Literal, Optional - -import grpc -import mujoco -import numpy as np -from numpy import typing as npt - -# INTERNAL IMPORT -from mujoco_mpc.proto import kalman_pb2 -from mujoco_mpc.proto import kalman_pb2_grpc - - -def find_free_port() -> int: - """Find an available TCP port on the system. - - This function creates a temporary socket, binds it to an available port - chosen by the operating system, and returns the chosen port number. - - Returns: - int: An available TCP port number. - """ - with socket.socket(family=socket.AF_INET6) as s: - s.bind(("", 0)) - s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - return s.getsockname()[1] - - -class Kalman: - """`Kalman` class to interface with MuJoCo MPC kalman. - - Attributes: - port: - channel: - stub: - server_process: - """ - - def __init__( - self, - model: mujoco.MjModel, - server_binary_path: Optional[str] = None, - send_as: Literal["mjb", "xml"] = "xml", - colab_logging: bool = True, - ): - # server - if server_binary_path is None: - binary_name = "kalman_server" - server_binary_path = pathlib.Path(__file__).parent / "mjpc" / binary_name - self._colab_logging = colab_logging - self.port = find_free_port() - self.server_process = subprocess.Popen( - [str(server_binary_path), f"--mjpc_port={self.port}"], - stdout=subprocess.PIPE if colab_logging else None, - ) - # os.set_blocking(self.server_process.stdout.fileno(), False) - atexit.register(self.server_process.kill) - - credentials = grpc.local_channel_credentials( - grpc.LocalConnectionType.LOCAL_TCP - ) - self.channel = grpc.secure_channel(f"localhost:{self.port}", credentials) - grpc.channel_ready_future(self.channel).result(timeout=10) - self.stub = kalman_pb2_grpc.KalmanStub(self.channel) - - # initialize - self.init( - model, - send_as=send_as, - ) - - def close(self): - self.channel.close() - self.server_process.kill() - self.server_process.wait() - - def init( - self, - model: mujoco.MjModel, - send_as: Literal["mjb", "xml"] = "xml", - ): - """ - Args: - model: optional `MjModel` instance, which, if provided, will be used as - the underlying model for planning. If not provided, the default MJPC - task xml will be used. - configuration_length: estimation horizon. - send_as: The serialization format for sending the model over gRPC; "xml". - """ - - # setup model - def model_to_mjb(model: mujoco.MjModel) -> bytes: - buffer_size = mujoco.mj_sizeModel(model) - buffer = np.empty(shape=buffer_size, dtype=np.uint8) - mujoco.mj_saveModel(model, None, buffer) - return buffer.tobytes() - - def model_to_xml(model: mujoco.MjModel) -> str: - tmp = tempfile.NamedTemporaryFile() - mujoco.mj_saveLastXML(tmp.name, model) - with pathlib.Path(tmp.name).open("rt") as f: - xml_string = f.read() - return xml_string - - if model is not None: - if send_as == "mjb": - model_message = kalman_pb2.MjModel(mjb=model_to_mjb(model)) - else: - model_message = kalman_pb2.MjModel(xml=model_to_xml(model)) - else: - model_message = None - - # initialize request - init_request = kalman_pb2.InitRequest( - model=model_message, - ) - - # initialize response - self._wait(self.stub.Init.future(init_request)) - - def reset(self): - # reset request - request = kalman_pb2.ResetRequest() - - # reset response - self._wait(self.stub.Reset.future(request)) - - def settings( - self, - epsilon: Optional[float] = None, - flg_centered: Optional[bool] = None, - ) -> dict[str, int | bool]: - # assemble settings - inputs = kalman_pb2.Settings( - epsilon=epsilon, - flg_centered=flg_centered, - ) - - # settings request - request = kalman_pb2.SettingsRequest( - settings=inputs, - ) - - # settings response - settings = self._wait(self.stub.Settings.future(request)).settings - - # return all settings - return { - "epsilon": settings.epsilon, - "flg_centered": settings.flg_centered, - } - - def update_measurement( - self, - ctrl: Optional[npt.ArrayLike] = [], - sensor: Optional[npt.ArrayLike] = [], - ): - # request - request = kalman_pb2.UpdateMeasurementRequest( - ctrl=ctrl, - sensor=sensor, - ) - - # response - self._wait(self.stub.UpdateMeasurement.future(request)) - - def update_prediction(self): - # request - request = kalman_pb2.UpdatePredictionRequest() - - # response - self._wait(self.stub.UpdatePrediction.future(request)) - - def timers(self) -> dict[str, float]: - # request - request = kalman_pb2.TimersRequest() - - # response - response = self._wait(self.stub.Timers.future(request)) - - # timers - return { - "measurement": response.measurement, - "prediction": response.prediction, - } - - def state(self, state: Optional[npt.ArrayLike] = []) -> np.ndarray: - # input - input = kalman_pb2.State(state=state) - - # request - request = kalman_pb2.StateRequest( - state=input, - ) - - # response - response = self._wait(self.stub.State.future(request)).state - - # return state - return np.array(response.state) - - def covariance( - self, covariance: Optional[npt.ArrayLike] = None - ) -> np.ndarray: - # input - inputs = kalman_pb2.Covariance( - covariance=covariance.flatten() if covariance is not None else None, - ) - - # request - request = kalman_pb2.CovarianceRequest( - covariance=inputs, - ) - - # response - response = self._wait(self.stub.Covariance.future(request)).covariance - - # return covariance - return np.array(response.covariance).reshape( - response.dimension, response.dimension - ) - - def noise( - self, - process: Optional[npt.ArrayLike] = [], - sensor: Optional[npt.ArrayLike] = [], - ) -> dict[str, np.ndarray]: - # inputs - inputs = kalman_pb2.Noise( - process=process, - sensor=sensor, - ) - - # request - request = kalman_pb2.NoiseRequest( - noise=inputs, - ) - - # response - response = self._wait(self.stub.Noise.future(request)).noise - - # return noise - return { - "process": np.array(response.process), - "sensor": np.array(response.sensor), - } - - def _wait(self, future): - """Waits for the future to complete, while printing out subprocess stdout.""" - # if self._colab_logging: - # while True: - # line = self.server_process.stdout.readline() - # if line: - # sys.stdout.write(line.decode("utf-8")) - # if future.done(): - # break - return future.result() diff --git a/python/mujoco_mpc/kalman_test.py b/python/mujoco_mpc/kalman_test.py deleted file mode 100644 index a95b5ef2e..000000000 --- a/python/mujoco_mpc/kalman_test.py +++ /dev/null @@ -1,103 +0,0 @@ -# Copyright 2023 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -from absl.testing import absltest -import mujoco -from mujoco_mpc import kalman as kalman_lib -import numpy as np - -import pathlib - - -class KalmanTest(absltest.TestCase): - - def test_settings(self): - # load model - model_path = ( - pathlib.Path(__file__).parent.parent.parent - / "mjpc/test/testdata/estimator/particle/task1D.xml" - ) - model = mujoco.MjModel.from_xml_path(str(model_path)) - - # initialize - kalman = kalman_lib.Kalman(model=model) - - # settings - epsilon = 2.0 - flg_centered = True - settings = kalman.settings(epsilon=epsilon, flg_centered=flg_centered) - - # test - self.assertLess(np.abs(settings["epsilon"] - epsilon), 1.0e-5) - self.assertTrue(settings["flg_centered"] == flg_centered) - - def test_updates(self): - # load model - model_path = ( - pathlib.Path(__file__).parent.parent.parent - / "mjpc/test/testdata/estimator/particle/task1D.xml" - ) - model = mujoco.MjModel.from_xml_path(str(model_path)) - - # initialize - kalman = kalman_lib.Kalman(model=model) - - # state - state = np.random.normal(scale=1.0, size=(model.nq + model.nv)) - state_response = kalman.state(state=state) - - # test state - self.assertLess(np.linalg.norm(state_response - state), 1.0e-5) - - # covariance - nvelocity = 2 * model.nv - F = np.random.normal(scale=1.0, size=(nvelocity**2)).reshape( - nvelocity, nvelocity - ) - covariance = F.T @ F - covariance_response = kalman.covariance(covariance=covariance) - - # test covariance - self.assertLess( - np.linalg.norm((covariance_response - covariance).ravel()), 1.0e-5 - ) - self.assertTrue(covariance_response.shape == (nvelocity, nvelocity)) - - # noise - process = np.random.normal(scale=1.0e-3, size=nvelocity) - sensor = np.random.normal(scale=1.0e-3, size=model.nsensordata) - noise = kalman.noise(process=process, sensor=sensor) - - # test noise - self.assertLess(np.linalg.norm(noise["process"] - process), 1.0e-5) - self.assertLess(np.linalg.norm(noise["sensor"] - sensor), 1.0e-5) - - # measurement update - ctrl = np.random.normal(scale=1.0, size=model.nu) - sensor = np.random.normal(scale=1.0, size=model.nsensordata) - kalman.update_measurement(ctrl=ctrl, sensor=sensor) - - # # prediction update - kalman.update_prediction() - - # timers - timer = kalman.timers() - - self.assertLess(0.0, timer["measurement"]) - self.assertLess(0.0, timer["prediction"]) - - -if __name__ == "__main__": - absltest.main() diff --git a/python/setup_kalman.py b/python/setup_kalman.py deleted file mode 100644 index 14a42e979..000000000 --- a/python/setup_kalman.py +++ /dev/null @@ -1,312 +0,0 @@ -# Copyright 2023 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Install script for MuJoCo MPC.""" - -import os -import pathlib -import platform -import shutil -import setuptools -from setuptools.command import build_ext -from setuptools.command import build_py -import subprocess - -Path = pathlib.Path - - -class GenerateProtoGrpcCommand(setuptools.Command): - """Specialized setup command to handle kalman proto compilation. - - Generates the `kalman_pb2{_grpc}.py` files from `kalman_proto`. Assumes that - `grpc_tools.protoc` is installed. - """ - - description = "Generate `.proto` files to Python protobuf and gRPC files." - user_options = [] - - def initialize_options(self): - self.build_lib = None - - def finalize_options(self): - self.set_undefined_options("build_py", ("build_lib", "build_lib")) - - def run(self): - """Generate `kalman.proto` into `kalman_pb2{_grpc}.py`. - - This function looks more complicated than what it has to be because the - `protoc` generator is very particular in the way it generates the imports - for the generated `kalman_pb2_grpc.py` file. The final argument of the - `protoc` call has to be "mujoco_mpc/kalman.proto" in order for the import to - become `from mujoco_mpc import [kalman_pb2_proto_import]` instead of just - `import [kalman_pb2_proto_import]`. The latter would fail because the name is - meant to be relative but python3 interprets it as an absolute import. - """ - # We import here because, if the import is at the top of this file, we - # cannot resolve the dependencies without having `grpcio-tools` installed. - from grpc_tools import protoc # pylint: disable=import-outside-toplevel - - kalman_proto_filename = "kalman.proto" - kalman_proto_source_path = Path( - "..", "grpc", kalman_proto_filename - ).resolve() - assert self.build_lib is not None - build_lib_path = Path(self.build_lib).resolve() - proto_module_relative_path = Path( - "mujoco_mpc", "proto", kalman_proto_filename - ) - kalman_proto_destination_path = Path( - build_lib_path, proto_module_relative_path - ) - kalman_proto_destination_path.parent.mkdir(parents=True, exist_ok=True) - # Copy `kalman_proto_filename` into current source. - shutil.copy(kalman_proto_source_path, kalman_proto_destination_path) - - protoc_command_parts = [ - # We use `__file__` as the first argument the same way as is done by - # `protoc` when called as `__main__` here: - # https://github.com/grpc/grpc/blob/21996c37842035661323c71b9e7040345f0915e2/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py#L172-L173. - __file__, - f"-I{build_lib_path}", - f"--python_out={build_lib_path}", - f"--grpc_python_out={build_lib_path}", - str(kalman_proto_destination_path), - ] - - protoc_returncode = protoc.main(protoc_command_parts) - - if protoc_returncode != 0: - raise subprocess.CalledProcessError( - returncode=protoc_returncode, - cmd=f"`protoc.main({protoc_command_parts})`", - ) - - self.spawn( - ["touch", str(kalman_proto_destination_path.parent / "__init__.py")] - ) - - -class CopyKalmanServerBinaryCommand(setuptools.Command): - """Specialized setup command to copy `kalman_server` next to `kalman.py`. - - Assumes that the C++ gRPC `kalman_server` binary has been manually built and - and located in the default `mujoco_mpc/build/bin` folder. - """ - - description = "Copy `kalman_server` next to `kalman.py`." - user_options = [] - - def initialize_options(self): - self.build_lib = None - - def finalize_options(self): - self.set_undefined_options("build_py", ("build_lib", "build_lib")) - - def run(self): - self._copy_binary("kalman_server") - # self._copy_binary("ui_kalman_server") - - def _copy_binary(self, binary_name): - source_path = Path(f"../build/bin/{binary_name}") - if not source_path.exists(): - raise ValueError( - f"Cannot find `{binary_name}` binary from {source_path}. Please build" - " the `{binary_name}` C++ gRPC service." - ) - assert self.build_lib is not None - build_lib_path = Path(self.build_lib).resolve() - destination_path = Path(build_lib_path, "mujoco_mpc", "mjpc", binary_name) - - self.announce(f"{source_path.resolve()=}") - self.announce(f"{destination_path.resolve()=}") - - destination_path.parent.mkdir(exist_ok=True, parents=True) - shutil.copy(source_path, destination_path) - - -class CopyTaskAssetsCommand(setuptools.Command): - """Copies `kalman_server` and `ui_kalman_server` next to `kalman.py`. - - Assumes that the C++ gRPC `kalman_server` binary has been manually built and - and located in the default `mujoco_mpc/build/bin` folder. - """ - - description = ( - "Copy task assets over to python source to make them accessible by" - " `Kalman`." - ) - user_options = [] - - def initialize_options(self): - self.build_lib = None - - def finalize_options(self): - self.set_undefined_options("build_ext", ("build_lib", "build_lib")) - - def run(self): - mjpc_tasks_path = Path(__file__).parent.parent / "mjpc" / "tasks" - source_paths = tuple(mjpc_tasks_path.rglob("*.xml")) - relative_source_paths = tuple( - p.relative_to(mjpc_tasks_path) for p in source_paths - ) - assert self.build_lib is not None - build_lib_path = Path(self.build_lib).resolve() - destination_dir_path = Path(build_lib_path, "mujoco_mpc", "mjpc", "tasks") - self.announce( - f"Copying assets {relative_source_paths} from" - f" {mjpc_tasks_path} over to {destination_dir_path}." - ) - - for source_path, relative_source_path in zip( - source_paths, relative_source_paths - ): - destination_path = destination_dir_path / relative_source_path - destination_path.parent.mkdir(exist_ok=True, parents=True) - shutil.copy(source_path, destination_path) - - -class BuildPyCommand(build_py.build_py): - """Specialized Python builder to handle kalman service dependencies. - - During build, this will generate the `kalman_pb2{_grpc}.py` files and copy - `kalman_server` binary next to `kalman.py`. - """ - - user_options = build_py.build_py.user_options - - def run(self): - self.run_command("generate_proto_grpc") - self.run_command("copy_task_assets") - super().run() - - -class CMakeExtension(setuptools.Extension): - """A Python extension that has been prebuilt by CMake. - - We do not want distutils to handle the build process for our extensions, so - so we pass an empty list to the super constructor. - """ - - def __init__(self, name): - super().__init__(name, sources=[]) - - -class BuildCMakeExtension(build_ext.build_ext): - """Uses CMake to build extensions.""" - - def run(self): - self._configure_and_build_kalman_server() - self.run_command("copy_kalman_server_binary") - - def _configure_and_build_kalman_server(self): - """Check for CMake.""" - cmake_command = "cmake" - build_cfg = "Debug" - mujoco_mpc_root = Path(__file__).parent.parent - mujoco_mpc_build_dir = mujoco_mpc_root / "build" - cmake_configure_args = [ - "-DCMAKE_EXPORT_COMPILE_COMMANDS:BOOL=TRUE", - f"-DCMAKE_BUILD_TYPE:STRING={build_cfg}", - "-DBUILD_TESTING:BOOL=OFF", - "-DMJPC_BUILD_GRPC_SERVICE:BOOL=ON", - ] - - if platform.system() == "Darwin" and "ARCHFLAGS" in os.environ: - osx_archs = [] - if "-arch x86_64" in os.environ["ARCHFLAGS"]: - osx_archs.append("x86_64") - if "-arch arm64" in os.environ["ARCHFLAGS"]: - osx_archs.append("arm64") - cmake_configure_args.append( - f"-DCMAKE_OSX_ARCHITECTURES={';'.join(osx_archs)}" - ) - - # TODO(hartikainen): We currently configure the builds into - # `mujoco_mpc/build`. This should use `self.build_{temp,lib}` instead, to - # isolate the Python builds from the C++ builds. - print("Configuring CMake with the following arguments:") - for arg in cmake_configure_args: - print(f" {arg}") - subprocess.check_call( - [ - cmake_command, - *cmake_configure_args, - f"-S{mujoco_mpc_root.resolve()}", - f"-B{mujoco_mpc_build_dir.resolve()}", - ], - cwd=mujoco_mpc_root, - ) - - print("Building `kalman_server` and `ui_kalman_server` with CMake") - subprocess.check_call( - [ - cmake_command, - "--build", - str(mujoco_mpc_build_dir.resolve()), - "--target", - "kalman_server", - # "ui_kalman_server", - f"-j{os.cpu_count()}", - "--config", - build_cfg, - ], - cwd=mujoco_mpc_root, - ) - - -setuptools.setup( - name="mujoco_mpc", - version="0.1.0", - author="DeepMind", - author_email="mujoco@deepmind.com", - description="MuJoCo MPC (MJPC)", - url="https://github.com/deepmind/mujoco_mpc", - license="MIT", - classifiers=[ - "Development Status :: 2 - Pre-Alpha", - "Intended Audience :: Developers", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: Apache Software License", - "Natural Language :: English", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3 :: Only", - "Topic :: Scientific/Engineering", - ], - packages=setuptools.find_packages(), - python_requires=">=3.7", - install_requires=[ - "grpcio-tools", - "grpcio", - ], - extras_require={ - "test": [ - "absl-py", - "mujoco >= 2.3.3", - ], - }, - ext_modules=[CMakeExtension("kalman_server")], - cmdclass={ - "build_py": BuildPyCommand, - "build_ext": BuildCMakeExtension, - "generate_proto_grpc": GenerateProtoGrpcCommand, - "copy_kalman_server_binary": CopyKalmanServerBinaryCommand, - "copy_task_assets": CopyTaskAssetsCommand, - }, - package_data={ - "": [ - "mjpc/kalman_server", - # "mjpc/ui_kalman_server", - ], - }, -)