Merge pull request #279 from google-deepmind/deepmind

Merge deepmind branch into main.

Author: nimrod-gileadi

.github/workflows/build.yml

@@ -68,7 +68,7 @@ jobs:
         $cmake_extra_args
     - name: Build MuJoCo MPC
       working-directory: build
-      run: cmake --build . --config=Release ${{ matrix.cmake_build_args }} --target mjpc agent_test agent_utilities_test cost_derivatives_test norm_test rollout_test threadpool_test trajectory_test direct_force_test direct_optimize_test direct_parameter_test direct_sensor_test direct_trajectory_test direct_utilities_test batch_filter_test batch_prior_test kalman_test unscented_test cubic_test gradient_planner_test gradient_test linear_test zero_test backward_pass_test ilqg_test robust_planner_test sampling_planner_test state_test task_test utilities_test ${{ matrix.additional_targets }}
+      run: cmake --build . --config=Release ${{ matrix.cmake_build_args }} --target mjpc agent_test agent_utilities_test cost_derivatives_test norm_test rollout_test threadpool_test trajectory_test direct_force_test direct_optimize_test direct_parameter_test direct_sensor_test direct_trajectory_test direct_utilities_test batch_filter_test batch_prior_test kalman_test spline_test unscented_test cubic_test gradient_planner_test gradient_test linear_test zero_test backward_pass_test ilqg_test robust_planner_test sampling_planner_test state_test task_test utilities_test ${{ matrix.additional_targets }}
     - name: Test MuJoCo MPC
       working-directory: build
       run: ctest -C Release --output-on-failure .

CMakeLists.txt

@@ -77,19 +77,12 @@ findorfetch(
   EXCLUDE_FROM_ALL
 )
 
+# TODO(nimrod): Update to the latest version of abseil, or use the one defined
+# by MuJoCo, once grpc fix their build issues.
 set(MUJOCO_DEP_VERSION_abseil
-    c8a2f92586fe9b4e1aff049108f5db8064924d8e # LTS 20230125.1
+    fb3621f4f897824c0dbe0615fa94543df6192f30 # LTS 20230802.1
     CACHE STRING "Version of `abseil` to be fetched."
-)
-
-set(MUJOCO_DEP_VERSION_glfw3
-    7482de6071d21db77a7236155da44c172a7f6c9e # 3.3.8
-    CACHE STRING "Version of `glfw` to be fetched."
-)
-
-set(MJPC_DEP_VERSION_lodepng
-    b4ed2cd7ecf61d29076169b49199371456d4f90b
-    CACHE STRING "Version of `lodepng` to be fetched."
+    FORCE
 )
 
 set(BUILD_SHARED_LIBS_OLD ${BUILD_SHARED_LIBS})
@@ -118,6 +111,9 @@ findorfetch(
 
 set(ABSL_PROPAGATE_CXX_STD ON)
 set(ABSL_BUILD_TESTING OFF)
+# ABSL_ENABLE_INSTALL is needed for
+# https://github.com/protocolbuffers/protobuf/issues/12185#issuecomment-1594685860
+set(ABSL_ENABLE_INSTALL ON)
 findorfetch(
   USE_SYSTEM_PACKAGE
   OFF

mjpc/CMakeLists.txt

@@ -38,6 +38,8 @@ add_library(
   tasks/tasks.h
   tasks/acrobot/acrobot.cc
   tasks/acrobot/acrobot.h
+  tasks/bimanual/bimanual.cc
+  tasks/bimanual/bimanual.h
   tasks/cartpole/cartpole.cc
   tasks/cartpole/cartpole.h
   tasks/cube/solve.cc
@@ -122,6 +124,8 @@ add_library(
   direct/trajectory.h
   direct/model_parameters.cc
   direct/model_parameters.h
+  spline/spline.cc
+  spline/spline.h
   app.cc
   app.h
   norm.cc
@@ -138,8 +142,11 @@ target_compile_definitions(libmjpc PRIVATE MJSIMULATE_STATIC)
 target_link_libraries(
   libmjpc
   absl::any_invocable
+  absl::check
   absl::flat_hash_map
+  absl::log
   absl::random_random
+  absl::span
   glfw
   lodepng
   mujoco::mujoco

mjpc/grpc/CMakeLists.txt

@@ -30,7 +30,7 @@ findorfetch(
   GIT_REPO
   https://github.com/grpc/grpc
   GIT_TAG
-  v1.53.0
+  v1.60.1
   TARGETS
   gRPC
 )

mjpc/grpc/agent.proto

@@ -39,6 +39,8 @@ service Agent {
       returns (GetTaskParametersResponse);
   // Set cost weights.
   rpc SetCostWeights(SetCostWeightsRequest) returns (SetCostWeightsResponse);
+  // Get cost term residuals.
+  rpc GetResiduals(GetResidualsRequest) returns (GetResidualsResponse);
   // Get cost term values.
   rpc GetCostValuesAndWeights(GetCostValuesAndWeightsRequest)
       returns (GetCostValuesAndWeightsResponse);
@@ -113,6 +115,16 @@ message GetActionResponse {
   repeated float action = 1 [packed = true];
 }
 
+message GetResidualsRequest {}
+
+message Residual {
+  repeated double values = 1;
+}
+
+message GetResidualsResponse {
+  map<string, Residual> values = 1;
+}
+
 message GetCostValuesAndWeightsRequest {}
 
 message ValueAndWeight {

mjpc/grpc/agent_service.cc

@@ -35,6 +35,8 @@ using ::agent::GetAllModesRequest;
 using ::agent::GetAllModesResponse;
 using ::agent::GetBestTrajectoryRequest;
 using ::agent::GetBestTrajectoryResponse;
+using ::agent::GetResidualsRequest;
+using ::agent::GetResidualsResponse;
 using ::agent::GetCostValuesAndWeightsRequest;
 using ::agent::GetCostValuesAndWeightsResponse;
 using ::agent::GetModeRequest;
@@ -181,6 +183,16 @@ grpc::Status AgentService::GetAction(grpc::ServerContext* context,
       request, &agent_, model, rollout_data_.get(), &rollout_state_, response);
 }
 
+grpc::Status AgentService::GetResiduals(
+    grpc::ServerContext* context, const GetResidualsRequest* request,
+    GetResidualsResponse* response) {
+  if (!Initialized()) {
+    return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."};
+  }
+  return grpc_agent_util::GetResiduals(request, &agent_, model,
+                                       data_, response);
+}
+
 grpc::Status AgentService::GetCostValuesAndWeights(
     grpc::ServerContext* context, const GetCostValuesAndWeightsRequest* request,
     GetCostValuesAndWeightsResponse* response) {

mjpc/grpc/agent_service.h

@@ -58,6 +58,11 @@ class AgentService final : public agent::Agent::Service {
                          const agent::GetActionRequest* request,
                          agent::GetActionResponse* response) override;
 
+  grpc::Status GetResiduals(
+      grpc::ServerContext* context,
+      const agent::GetResidualsRequest* request,
+      agent::GetResidualsResponse* response) override;
+
   grpc::Status GetCostValuesAndWeights(
       grpc::ServerContext* context,
       const agent::GetCostValuesAndWeightsRequest* request,

mjpc/grpc/agent_service_test.cc

@@ -360,4 +360,16 @@ TEST_F(AgentServiceTest, GetAllModes_Works) {
   EXPECT_EQ(response.mode_names()[0], "default_mode");
 }
 
+TEST_F(AgentServiceTest, GetResiduals_Works) {
+  RunAndCheckInit("Cartpole", nullptr);
+
+  grpc::ClientContext context;
+
+  agent::GetResidualsRequest request;
+  agent::GetResidualsResponse response;
+  grpc::Status status = stub->GetResiduals(&context, request, &response);
+
+  EXPECT_TRUE(status.ok());
+}
+
 }  // namespace mjpc::agent_grpc

mjpc/grpc/grpc_agent_util.cc

@@ -44,6 +44,8 @@ using ::agent::GetActionRequest;
 using ::agent::GetActionResponse;
 using ::agent::GetAllModesRequest;
 using ::agent::GetAllModesResponse;
+using ::agent::GetResidualsRequest;
+using ::agent::GetResidualsResponse;
 using ::agent::GetCostValuesAndWeightsRequest;
 using ::agent::GetCostValuesAndWeightsResponse;
 using ::agent::GetModeRequest;
@@ -58,6 +60,7 @@ using ::agent::SetCostWeightsRequest;
 using ::agent::SetModeRequest;
 using ::agent::SetStateRequest;
 using ::agent::SetTaskParametersRequest;
+using ::agent::Residual;
 using ::agent::ValueAndWeight;
 
 grpc::Status GetState(const mjModel* model, const mjData* data,
@@ -226,6 +229,34 @@ grpc::Status GetAction(const GetActionRequest* request,
   return grpc::Status::OK;
 }
 
+grpc::Status GetResiduals(
+    const GetResidualsRequest* request, const mjpc::Agent* agent,
+    const mjModel* model, mjData* data,
+    GetResidualsResponse* response) {
+  const mjModel* agent_model = agent->GetModel();
+  const mjpc::Task* task = agent->ActiveTask();
+  std::vector<double> residuals(task->num_residual, 0);  // scratch space
+  task->Residual(model, data, residuals.data());
+  std::vector<int> dim_norm_residual = task->dim_norm_residual;
+
+  int residual_shift = 0;
+  for (int i = 0; i < task->num_term; i++) {
+    CHECK_EQ(agent_model->sensor_type[i], mjSENS_USER);
+    std::string_view sensor_name(agent_model->names +
+                                  agent_model->name_sensoradr[i]);
+
+    std::vector<double> sensor_residual_values(
+        residuals.begin() + residual_shift,
+        residuals.begin() + residual_shift + dim_norm_residual[i]);
+    Residual sensor_residual;
+    sensor_residual.mutable_values()->Assign(sensor_residual_values.begin(),
+                                             sensor_residual_values.end());
+    (*response->mutable_values())[sensor_name] = sensor_residual;
+    residual_shift += dim_norm_residual[i];
+  }
+  return grpc::Status::OK;
+}
+
 grpc::Status GetCostValuesAndWeights(
     const GetCostValuesAndWeightsRequest* request, const mjpc::Agent* agent,
     const mjModel* model, mjData* data,

mjpc/grpc/grpc_agent_util.h

@@ -34,6 +34,10 @@ grpc::Status GetAction(const agent::GetActionRequest* request,
                        const mjModel* model, mjData* rollout_data,
                        mjpc::State* rollout_state,
                        agent::GetActionResponse* response);
+grpc::Status GetResiduals(
+    const agent::GetResidualsRequest* request,
+    const mjpc::Agent* agent, const mjModel* model, mjData* data,
+    agent::GetResidualsResponse* response);
 grpc::Status GetCostValuesAndWeights(
     const agent::GetCostValuesAndWeightsRequest* request,
     const mjpc::Agent* agent, const mjModel* model, mjData* data,