Merge branch 'main' into improve_docs

google-deepmind · Jan 12, 2024 · 4572f3d · 4572f3d
2 parents 05d292f + fd9c479
commit 4572f3d
Show file tree

Hide file tree

Showing 42 changed files with 472 additions and 108 deletions.
diff --git a/README.md b/README.md
@@ -42,10 +42,30 @@ For a quick video overview of MJPC, click below.
 
 [![Video](http://img.youtube.com/vi/Bdx7DuAMB6o/hqdefault.jpg)](https://dpmd.ai/mjpc)
 
-For a longer talk at the MIT Robotics Seminar describing our results, click
+For a longer talk at the MIT Robotics Seminar in December 2022 describing our results, click
 below.
 
-[![Talk](http://img.youtube.com/vi/2xVN-qY78P4/hqdefault.jpg)](https://www.youtube.com/watch?v=2xVN-qY78P4)
+[![2022Talk](http://img.youtube.com/vi/2xVN-qY78P4/hqdefault.jpg)](https://www.youtube.com/watch?v=2xVN-qY78P4)
+
+A more recent, December 2023 talk at the IEEE Technical Committee on Model-Based Optimization
+is available here:
+
+[![2023Talk](https://img.youtube.com/vi/J-JO-lgaKtw/hqdefault.jpg)](https://www.youtube.com/watch?v=J-JO-lgaKtw&t=0s)
+
+### Example tasks
+
+Quadruped task:
+
+[![Quadruped](http://img.youtube.com/vi/esLuwaWz4oE/hqdefault.jpg)](https://www.youtube.com/watch?v=esLuwaWz4oE)
+
+Rubik's cube 10-move unscramble:
+
+[![Unscramble](http://img.youtube.com/vi/ZRRvVWV-Muk/hqdefault.jpg)](https://www.youtube.com/watch?v=ZRRvVWV-Muk)
+
+Humanoid motion-capture tracking:
+
+[![Tracking](http://img.youtube.com/vi/tEBVK-MO1Sw/hqdefault.jpg)](https://www.youtube.com/watch?v=tEBVK-MO1Sw)
+
 
 ## Graphical User Interface
 
@@ -81,7 +101,7 @@ Note, we are using `clang-14`.
 
 # Python API
 
-We provide a simple Python API for MJPC. This API is still experimental and expects some more experience from its users. For example, the correct usage requires that the model (defined in Python) and the MJPC task (i.e., the residual and transition functions defined in C++) are compatible with each other. Currently, the Python API does not provide any particular error handling for verifying this compatibilty and may be difficult to debug without more in-depth knowedge about MuJoCo and MJPC.
+We provide a simple Python API for MJPC. This API is still experimental and expects some more experience from its users. For example, the correct usage requires that the model (defined in Python) and the MJPC task (i.e., the residual and transition functions defined in C++) are compatible with each other. Currently, the Python API does not provide any particular error handling for verifying this compatibility and may be difficult to debug without more in-depth knowledge about mujoco and MJPC.
 
 - [agent.py](python/mujoco_mpc/agent.py) for available methods for planning.
 

diff --git a/docs/OVERVIEW.md b/docs/OVERVIEW.md
@@ -217,7 +217,7 @@ As an example, consider these snippets specifying the Swimmer task:
 </sensor>
 ```
 ```c++
-void Swimmer::ResidualFn::Residual(const mjModel* model, 
+void Swimmer::ResidualFn::Residual(const mjModel* model,
                                    const mjData* data,
                                    double* residual) const {
   // initialize counter

diff --git a/mjpc/agent.cc b/mjpc/agent.cc
@@ -159,10 +159,10 @@ void Agent::Allocate() {
 }
 
 // reset data, settings, planners, state
-void Agent::Reset() {
+void Agent::Reset(const double* initial_repeated_action) {
   // planner
   for (const auto& planner : planners_) {
-    planner->Reset(kMaxTrajectoryHorizon);
+    planner->Reset(kMaxTrajectoryHorizon, initial_repeated_action);
   }
 
   // state

diff --git a/mjpc/agent.h b/mjpc/agent.h
@@ -65,7 +65,7 @@ class Agent {
   void Allocate();
 
   // reset data, settings, planners, states
-  void Reset();
+  void Reset(const double* initial_repeated_action = nullptr);
 
   // single planner iteration
   void PlanIteration(ThreadPool* pool);

diff --git a/mjpc/app.cc b/mjpc/app.cc
@@ -232,12 +232,16 @@ void PhysicsLoop(mj::Simulate& sim) {
         sim.agent->plot_enabled = absl::GetFlag(FLAGS_show_plot);
         sim.agent->plan_enabled = absl::GetFlag(FLAGS_planner_enabled);
         sim.agent->Allocate();
-        sim.agent->Reset();
-        sim.agent->PlotInitialize();
 
         // set home keyframe
         int home_id = mj_name2id(mnew, mjOBJ_KEY, "home");
-        if (home_id >= 0) mj_resetDataKeyframe(mnew, dnew, home_id);
+        if (home_id >= 0) {
+          mj_resetDataKeyframe(mnew, dnew, home_id);
+          sim.agent->Reset(dnew->ctrl);
+        } else {
+          sim.agent->Reset();
+        }
+        sim.agent->PlotInitialize();
 
         sim.Load(mnew, dnew, sim.filename, true);
         m = mnew;
@@ -305,7 +309,7 @@ void PhysicsLoop(mj::Simulate& sim) {
           double slowdown = 100 / sim.percentRealTime[sim.real_time_index];
 
           // misalignment condition: distance from target sim time is bigger
-          // than syncmisalign
+          // than maximum misalignment `syncMisalign`
           bool misaligned = mju_abs(Seconds(elapsedCPU).count() / slowdown -
                                     elapsedSim) > syncMisalign;
 

diff --git a/mjpc/grpc/agent.proto b/mjpc/grpc/agent.proto
@@ -48,8 +48,13 @@ service Agent {
   rpc GetMode(GetModeRequest) returns (GetModeResponse);
   // Get all modes.
   rpc GetAllModes(GetAllModesRequest) returns (GetAllModesResponse);
+
   // Get best trajectory (states, actions, times).
-  rpc GetBestTrajectory(GetBestTrajectoryRequest) returns (GetBestTrajectoryResponse);
+  rpc GetBestTrajectory(GetBestTrajectoryRequest)
+      returns (GetBestTrajectoryResponse);
+
+  // A single method that can set many of the inputs.
+  rpc SetAnything(SetAnythingRequest) returns (SetAnythingResponse);
 }
 
 message MjModel {
@@ -186,3 +191,24 @@ message GetBestTrajectoryResponse {
   repeated double times = 3 [packed = true];
   int32 steps = 4;
 }
+
+message Pose {
+  repeated double pos = 1 [packed = true];
+  repeated double quat = 2 [packed = true];
+}
+
+message SetAnythingRequest {
+  State state = 1;
+
+  // map from parameter name to desired value
+  map<string, TaskParameterValue> parameters = 2;
+  // cost weights by name
+  map<string, double> cost_weights = 3;
+  string mode = 4;
+
+  // set the positions of mocap bodies by name. If `state` is set too, mocap
+  // positions will be set after the state is set.
+  map<string, Pose> mocap = 5;
+}
+
+message SetAnythingResponse {}
diff --git a/mjpc/grpc/agent_service.cc b/mjpc/grpc/agent_service.cc
@@ -25,6 +25,7 @@
 #include "mjpc/grpc/agent.pb.h"
 #include "mjpc/grpc/grpc_agent_util.h"
 #include "mjpc/task.h"
+#include "mjpc/trajectory.h"
 
 namespace mjpc::agent_grpc {
 
@@ -48,6 +49,8 @@ using ::agent::PlannerStepRequest;
 using ::agent::PlannerStepResponse;
 using ::agent::ResetRequest;
 using ::agent::ResetResponse;
+using ::agent::SetAnythingRequest;
+using ::agent::SetAnythingResponse;
 using ::agent::SetCostWeightsRequest;
 using ::agent::SetCostWeightsResponse;
 using ::agent::SetModeRequest;
@@ -323,4 +326,14 @@ grpc::Status AgentService::GetBestTrajectory(
   return grpc::Status::OK;
 }
 
+
+grpc::Status AgentService::SetAnything(
+    grpc::ServerContext* context, const SetAnythingRequest* request,
+    SetAnythingResponse* response) {
+  if (!Initialized()) {
+    return {grpc::StatusCode::FAILED_PRECONDITION, "Init not called."};
+  }
+  return grpc_agent_util::SetAnything(request, &agent_, agent_.GetModel(),
+                                      data_, response);
+}
 }  // namespace mjpc::agent_grpc
diff --git a/mjpc/grpc/agent_service.h b/mjpc/grpc/agent_service.h
@@ -109,6 +109,10 @@ class AgentService final : public agent::Agent::Service {
       const agent::GetBestTrajectoryRequest* request,
       agent::GetBestTrajectoryResponse* response) override;
 
+  grpc::Status SetAnything(grpc::ServerContext* context,
+                           const agent::SetAnythingRequest* request,
+                           agent::SetAnythingResponse* response) override;
+
  private:
   bool Initialized() const { return data_ != nullptr; }
 

diff --git a/mjpc/grpc/grpc_agent_util.cc b/mjpc/grpc/grpc_agent_util.cc
@@ -14,6 +14,7 @@
 
 #include "mjpc/grpc/grpc_agent_util.h"
 
+#include <cstring>
 #include <memory>
 #include <sstream>
 #include <string>
@@ -35,6 +36,7 @@
 #include "mjpc/agent.h"
 #include "mjpc/states/state.h"
 #include "mjpc/task.h"
+#include "mjpc/utilities.h"
 
 namespace grpc_agent_util {
 
@@ -49,6 +51,9 @@ using ::agent::GetModeResponse;
 using ::agent::GetStateResponse;
 using ::agent::GetTaskParametersRequest;
 using ::agent::GetTaskParametersResponse;
+using ::agent::Pose;
+using ::agent::SetAnythingRequest;
+using ::agent::SetAnythingResponse;
 using ::agent::SetCostWeightsRequest;
 using ::agent::SetModeRequest;
 using ::agent::SetStateRequest;
@@ -104,10 +109,9 @@ if (!(expr).ok()) { \
   } \
 }
 
-grpc::Status SetState(const SetStateRequest* request, mjpc::Agent* agent,
+namespace {
+grpc::Status SetState(const agent::State& state, mjpc::Agent* agent,
                       const mjModel* model, mjData* data) {
-  agent::State state = request->state();
-
   if (state.has_time()) data->time = state.time();
 
   if (state.qpos_size() > 0) {
@@ -144,6 +148,13 @@ grpc::Status SetState(const SetStateRequest* request, mjpc::Agent* agent,
 
   return grpc::Status::OK;
 }
+}  // namespace
+
+grpc::Status SetState(const SetStateRequest* request, mjpc::Agent* agent,
+                      const mjModel* model, mjData* data) {
+  agent::State state = request->state();
+  return SetState(state, agent, model, data);
+}
 
 #undef CHECK_SIZE
 
@@ -328,12 +339,10 @@ grpc::Status GetTaskParameters(const GetTaskParametersRequest* request,
   return grpc::Status::OK;
 }
 
-grpc::Status SetCostWeights(const SetCostWeightsRequest* request,
-                            mjpc::Agent* agent) {
-  if (request->reset_to_defaults()) {
-    agent->ActiveTask()->Reset(agent->GetModel());
-  }
-  for (const auto& [name, weight] : request->cost_weights()) {
+grpc::Status SetCostWeights(
+    const ::google::protobuf::Map<std::string, double>& cost_weights,
+    mjpc::Agent* agent) {
+  for (const auto& [name, weight] : cost_weights) {
     if (agent->SetWeightByName(name, weight) == -1) {
       std::ostringstream error_string;
       error_string << "Weight '" << name
@@ -346,19 +355,29 @@ grpc::Status SetCostWeights(const SetCostWeightsRequest* request,
                                      agent_model->name_sensoradr[i]);
         error_string << "  " << sensor_name << "\n";
       }
-      return {grpc::StatusCode::INVALID_ARGUMENT, error_string.str()};
+      return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
+                          error_string.str());
     }
   }
-
   return grpc::Status::OK;
 }
 
-grpc::Status SetMode(const SetModeRequest* request, mjpc::Agent* agent) {
-  int outcome = agent->SetModeByName(request->mode());
+grpc::Status SetCostWeights(const SetCostWeightsRequest* request,
+                            mjpc::Agent* agent) {
+  if (request->reset_to_defaults()) {
+    agent->ActiveTask()->Reset(agent->GetModel());
+  }
+  auto cost_weights = request->cost_weights();
+  return SetCostWeights(cost_weights, agent);
+}
+
+
+grpc::Status SetMode(std::string_view mode, mjpc::Agent* agent) {
+  int outcome = agent->SetModeByName(mode);
   if (outcome == -1) {
     std::vector<std::string> mode_names = agent->GetAllModeNames();
     std::ostringstream error_string;
-    error_string << "Mode '" << request->mode()
+    error_string << "Mode '" << mode
                   << "' not found in task. Available names are:\n";
     for (const auto& mode_name : mode_names) {
       error_string << "  " << mode_name << "\n";
@@ -369,6 +388,10 @@ grpc::Status SetMode(const SetModeRequest* request, mjpc::Agent* agent) {
   }
 }
 
+grpc::Status SetMode(const SetModeRequest* request, mjpc::Agent* agent) {
+  return SetMode(request->mode(), agent);
+}
+
 grpc::Status GetMode(const GetModeRequest* request, mjpc::Agent* agent,
                      GetModeResponse* response) {
   response->set_mode(agent->GetModeName());
@@ -384,6 +407,83 @@ grpc::Status GetAllModes(const GetAllModesRequest* request, mjpc::Agent* agent,
   return grpc::Status::OK;
 }
 
+namespace {
+grpc::Status SetMocap(const ::google::protobuf::Map<std::string, Pose>& mocap,
+                      mjpc::Agent* agent, const mjModel* model, mjData* data) {
+  // Check all names and poses before applying changes.
+  for (const auto& [name, pose] : mocap) {
+    int id = mj_name2id(model, mjOBJ_BODY, name.c_str());
+    if (id < 0) {
+      return grpc::Status(grpc::StatusCode::INVALID_ARGUMENT,
+                          absl::StrFormat("Body '%s' not found.", name));
+    }
+    int mocap_id = model->body_mocapid[id];
+    if (mocap_id < 0) {
+      return grpc::Status(
+          grpc::StatusCode::INVALID_ARGUMENT,
+          absl::StrFormat("Body '%s' is not a mocap body.", name));
+    }
+    if (pose.pos_size() != 0 && pose.pos_size() != 3) {
+      return grpc::Status(
+          grpc::StatusCode::INVALID_ARGUMENT,
+          absl::StrFormat("Mocap '%s' has invalid pose size %d.", name,
+                          pose.pos_size()));
+    }
+    if (pose.quat_size() != 0 && pose.quat_size() != 4) {
+      return grpc::Status(
+          grpc::StatusCode::INVALID_ARGUMENT,
+          absl::StrFormat("Mocap '%s' has invalid quat size %d.", name,
+                          pose.pos_size()));
+    }
+  }
+  for (const auto& [name, pose] : mocap) {
+    int id = mj_name2id(model, mjOBJ_BODY, name.c_str());
+    int mocap_id = model->body_mocapid[id];
+    for (int i = 0; i < pose.pos_size(); i++) {
+      data->mocap_pos[3*mocap_id + i] = pose.pos(i);
+    }
+    if (pose.quat_size() == 4) {
+      for (int i = 0; i < 4; i++) {
+        data->mocap_quat[4*mocap_id + i] = pose.quat(i);
+      }
+      mju_normalize4(data->mocap_quat + 4*mocap_id);
+    }
+  }
+  agent->SetState(data);
+  return grpc::Status::OK;
+}
+}  // namespace
+
+grpc::Status SetAnything(const SetAnythingRequest* request, mjpc::Agent* agent,
+                         const mjModel* model, mjData* data,
+                         SetAnythingResponse* response) {
+  if (request->has_state()) {
+    grpc::Status status = SetState(request->state(), agent, model, data);
+    if (!status.ok()) {
+      return status;
+    }
+  }
+  if (request->cost_weights_size() > 0) {
+    grpc::Status status = SetCostWeights(request->cost_weights(), agent);
+    if (!status.ok()) {
+      return status;
+    }
+  }
+  if (!request->mode().empty()) {
+    grpc::Status status = SetMode(request->mode(), agent);
+    if (!status.ok()) {
+      return status;
+    }
+  }
+  if (request->mocap_size() > 0) {
+    grpc::Status status = SetMocap(request->mocap(), agent, model, data);
+    if (!status.ok()) {
+      return status;
+    }
+  }
+  return grpc::Status::OK;
+}
+
 mjpc::UniqueMjModel LoadModelFromString(std::string_view xml, char* error,
                              int error_size) {
   static constexpr char file[] = "temporary-filename.xml";
@@ -392,7 +492,7 @@ mjpc::UniqueMjModel LoadModelFromString(std::string_view xml, char* error,
   mj_defaultVFS(vfs.get());
   mj_makeEmptyFileVFS(vfs.get(), file, xml.size());
   int file_idx = mj_findFileVFS(vfs.get(), file);
-  memcpy(vfs->filedata[file_idx], xml.data(), xml.size());
+  std::memcpy(vfs->filedata[file_idx], xml.data(), xml.size());
   mjpc::UniqueMjModel m = {mj_loadXML(file, vfs.get(), error, error_size),
                            mj_deleteModel};
   mj_deleteFileVFS(vfs.get(), file);