Added save/load and sample motion strategies

CMakeLists.txt

@@ -212,5 +212,5 @@ endif()
 set(MJPC_COMPILE_OPTIONS "${AVX_COMPILE_OPTIONS}" "${EXTRA_COMPILE_OPTIONS}")
 set(MJPC_LINK_OPTIONS "${EXTRA_LINK_OPTIONS}")
 
+add_definitions(-DSOURCE_DIR=\"${CMAKE_SOURCE_DIR}\")
 add_subdirectory(mjpc)
-

mjpc/CMakeLists.txt

@@ -165,6 +165,7 @@ target_link_libraries(
   mujoco::platform_ui_adapter
   threadpool
   Threads::Threads
+  nlohmann_json::nlohmann_json
 )
 target_include_directories(libmjpc
   PUBLIC
@@ -238,3 +239,7 @@ endif()
 if(MJPC_BUILD_GRPC_SERVICE)
   add_subdirectory(grpc)
 endif()
+
+include(FetchContent)
+FetchContent_Declare(json URL https://github.com/nlohmann/json/releases/download/v3.11.3/json.tar.xz)
+FetchContent_MakeAvailable(json)

mjpc/task.cc

@@ -202,13 +202,15 @@ void Task::Reset(const mjModel* model) {
 
   // loop over sensors
   int parameter_shift = 0;
+  weight_names.resize(num_term);
   for (int i = 0; i < num_term; i++) {
     // residual dimension
     num_residual += model->sensor_dim[i];
     dim_norm_residual[i] = (int)model->sensor_dim[i];
 
     // user data: [norm, weight, weight_lower, weight_upper, parameters...]
     double* s = model->sensor_user + i * model->nuser_sensor;
+    char* name = model->names + model->name_sensoradr[i];
 
     // check number of parameters
     int norm_parameter_dimension = NormParameterDimension(s[0]);
@@ -231,6 +233,7 @@ void Task::Reset(const mjModel* model) {
     }
 
     weight[i] = s[1];
+    weight_names[i] = name;
     num_norm_parameter[i] = norm_parameter_dimension;
     mju_copy(DataAt(norm_parameter, parameter_shift), s + 4,
              num_norm_parameter[i]);

mjpc/task.h

@@ -139,6 +139,7 @@ class Task {
   std::vector<int> num_norm_parameter;
   std::vector<NormType> norm;
   std::vector<double> weight;
+  std::vector<std::string> weight_names;
   std::vector<double> norm_parameter;
   double risk;
 

mjpc/tasks/humanoid/interact/contact_keyframe.cc

@@ -46,4 +46,42 @@ void ContactKeyframe::Reset() {
   facing_target.clear();
   weight.clear();
 }
+
+void to_json(json& j, const ContactPair& contact_pair) {
+  j = json{{"body1", contact_pair.body1},
+           {"body2", contact_pair.body2},
+           {"geom1", contact_pair.geom1},
+           {"geom2", contact_pair.geom2},
+           {"local_pos1", contact_pair.local_pos1},
+           {"local_pos2", contact_pair.local_pos2}};
+}
+
+void from_json(const json& j, ContactPair& contact_pair) {
+  j.at("body1").get_to(contact_pair.body1);
+  j.at("body2").get_to(contact_pair.body2);
+  j.at("geom1").get_to(contact_pair.geom1);
+  j.at("geom2").get_to(contact_pair.geom2);
+  j.at("local_pos1").get_to(contact_pair.local_pos1);
+  j.at("local_pos2").get_to(contact_pair.local_pos2);
+}
+
+void to_json(json& j, const ContactKeyframe& keyframe) {
+  j = json{{"name", keyframe.name},
+           {"contacts", keyframe.contact_pairs},
+           {"facing_target", keyframe.facing_target},
+           {"time_limit", keyframe.time_limit},
+           {"success_sustain_time", keyframe.success_sustain_time},
+           {"target_distance_tolerance", keyframe.target_distance_tolerance},
+           {"weight", keyframe.weight}};
+}
+
+void from_json(const json& j, ContactKeyframe& keyframe) {
+  j.at("name").get_to(keyframe.name);
+  j.at("contacts").get_to(keyframe.contact_pairs);
+  j.at("facing_target").get_to(keyframe.facing_target);
+  j.at("time_limit").get_to(keyframe.time_limit);
+  j.at("success_sustain_time").get_to(keyframe.success_sustain_time);
+  j.at("target_distance_tolerance").get_to(keyframe.target_distance_tolerance);
+  j.at("weight").get_to(keyframe.weight);
+}
 }  // namespace mjpc::humanoid

mjpc/tasks/humanoid/interact/contact_keyframe.h

@@ -15,12 +15,15 @@
 #ifndef MJPC_TASKS_HUMANOID_INTERACT_CONTACT_KEYFRAME_H_
 #define MJPC_TASKS_HUMANOID_INTERACT_CONTACT_KEYFRAME_H_
 
-#include <mujoco/mujoco.h>
 
 #include <map>
 #include <string>
 #include <vector>
 
+#include <mujoco/mujoco.h>
+#include "nlohmann/json.hpp"
+using json = nlohmann::json;
+
 namespace mjpc::humanoid {
 
 // ---------- Constants ----------------- //
@@ -86,6 +89,11 @@ class ContactKeyframe {
                                      // certain time
 };
 
+void to_json(json& j, const ContactPair& contact_pair);
+void from_json(const json& j, ContactPair& contact_pair);
+void to_json(json& j, const ContactKeyframe& keyframe);
+void from_json(const json& j, ContactKeyframe& keyframe);
+
 }  // namespace mjpc::humanoid
 
 #endif  // MJPC_TASKS_HUMANOID_INTERACT_CONTACT_KEYFRAME_H_

mjpc/tasks/humanoid/interact/interact.cc

@@ -124,6 +124,39 @@ void Interact::ResidualFn::ContactResidual(const mjModel* model,
   }
 }
 
+void Interact::SaveParamsToKeyframe(ContactKeyframe& kf) const {
+  for (int i = 0; i < weight_names.size(); i++) {
+    kf.weight[weight_names[i]] = weight[i];
+  }
+  kf.target_distance_tolerance = parameters[kDistanceToleranceParameterIndex];
+  kf.time_limit = parameters[kTimeLimitParameterIndex];
+  kf.success_sustain_time = parameters[kSustainTimeParameterIndex];
+}
+
+void Interact::LoadParamsFromKeyframe(const ContactKeyframe& kf) {
+  ContactKeyframe current_keyframe = motion_strategy_.GetCurrentKeyframe();
+  weight.clear();
+  int index = 0;
+  for (auto& w : weight_names) {
+    if (kf.weight.find(w) != kf.weight.end()) {
+      weight.push_back(kf.weight.at(w));
+    } else {
+      double default_weight =
+          default_weights[residual_.current_task_mode_][index];
+      std::printf(
+          "Keyframe %s does not have weight for %s, set default %.1f.\n",
+          kf.name.c_str(), w.c_str(), default_weight);
+      weight.push_back(default_weight);
+    }
+    current_keyframe.weight[w] = weight[index];
+    index++;
+  }
+  current_keyframe.name = kf.name;
+  parameters[kDistanceToleranceParameterIndex] = kf.target_distance_tolerance;
+  parameters[kTimeLimitParameterIndex] = kf.time_limit;
+  parameters[kSustainTimeParameterIndex] = kf.success_sustain_time;
+}
+
 // ---------------- Residuals for humanoid interaction task ----------- //
 //   Number of residuals: 13
 //     Residual (0): Torso up
@@ -181,13 +214,18 @@ void Interact::ResidualFn::Residual(const mjModel* model, const mjData* data,
 
 // -------- Transition for interaction task -------- //
 void Interact::TransitionLocked(mjModel* model, mjData* data) {
-  //   If the task mode is changed, sync it with the residual here
+  //  If the task mode is changed, sync it with the residual here
   if (residual_.current_task_mode_ != mode) {
     residual_.current_task_mode_ = (TaskMode)mode;
     weight = default_weights[residual_.current_task_mode_];
   }
 
-  if (!motion_strategy_.HasKeyframes()) return;
+  //  If the motion strategy is not initialized, load the given strategy
+  if (!motion_strategy_.HasKeyframes()) {
+    motion_strategy_.LoadStrategy("armchair_cross_leg");
+    LoadParamsFromKeyframe(motion_strategy_.GetCurrentKeyframe());
+    return;
+  }
 
   const ContactKeyframe& current_keyframe =
       motion_strategy_.GetCurrentKeyframe();
@@ -197,8 +235,9 @@ void Interact::TransitionLocked(mjModel* model, mjData* data) {
   if (data->time - motion_strategy_.GetCurrentKeyframeStartTime() >
           current_keyframe.time_limit &&
       total_distance > current_keyframe.target_distance_tolerance) {
-    // success criteria not reached, reset
+    // timelimit reached but distance criteria not reached, reset
     motion_strategy_.Reset();
+    LoadParamsFromKeyframe(motion_strategy_.GetCurrentKeyframe());
     residual_.residual_keyframe_ = motion_strategy_.GetCurrentKeyframe();
     motion_strategy_.SetCurrentKeyframeStartTime(data->time);
   } else if (total_distance <= current_keyframe.target_distance_tolerance &&
@@ -207,11 +246,13 @@ void Interact::TransitionLocked(mjModel* model, mjData* data) {
                  current_keyframe.success_sustain_time) {
     // success criteria reached, go to the next keyframe
     motion_strategy_.NextKeyframe();
+    LoadParamsFromKeyframe(motion_strategy_.GetCurrentKeyframe());
     residual_.residual_keyframe_ = motion_strategy_.GetCurrentKeyframe();
   } else if (total_distance > current_keyframe.target_distance_tolerance) {
     // keyframe error is more than tolerance, update the success start time
     motion_strategy_.SetCurrentKeyframeSuccessStartTime(data->time);
   }
+  SaveParamsToKeyframe(motion_strategy_.GetCurrentKeyframe());
 }
 
 // draw task-related geometry in the scene

mjpc/tasks/humanoid/interact/interact.h

@@ -27,9 +27,12 @@
 namespace mjpc::humanoid {
 
 // ---------- Constants ----------------- //
+constexpr int kNumberOfFreeJoints = 0;
 constexpr int kHeadHeightParameterIndex = 0;
 constexpr int kTorsoHeightParameterIndex = 1;
-constexpr int kNumberOfFreeJoints = 0;
+constexpr int kDistanceToleranceParameterIndex = 2;
+constexpr int kTimeLimitParameterIndex = 3;
+constexpr int kSustainTimeParameterIndex = 4;
 
 // ---------- Enums ----------------- //
 enum TaskMode : int {
@@ -124,6 +127,9 @@ class Interact : public Task {
   ResidualFn residual_;
   MotionStrategy motion_strategy_;
 
+  void SaveParamsToKeyframe(ContactKeyframe& kf) const;
+  void LoadParamsFromKeyframe(const ContactKeyframe& kf);
+
   // draw task-related geometry in the scene
   void ModifyScene(const mjModel* model, const mjData* data,
                    mjvScene* scene) const override;

mjpc/tasks/humanoid/interact/motion_strategy.cc

@@ -78,4 +78,70 @@ double MotionStrategy::CalculateTotalKeyframeDistance(
   return total_error;
 }
 
+bool MotionStrategy::SaveStrategy(const std::string& name,
+                                  const std::string& path) {
+  std::string filename(path);
+  filename.append(name);
+  filename.append(".json");
+
+  try {
+    json json_keyframes = json::array();
+    json json_kf;
+    for (auto& kf: contact_keyframes_) {
+      to_json(json_kf, kf);
+      json_keyframes.push_back(json_kf);
+    }
+
+    std::ofstream of(filename);
+    if (!of.is_open()) {
+      std::printf("\nFailed to open file %s", filename.c_str());
+      return false;
+    }
+    of << std::setw(4) << json_keyframes << std::endl;
+    of.close();
+    std::printf("\nKeyframe sequence saved as %s", filename.c_str());
+    return true;
+  } catch (const std::exception &e) {
+    std::cerr << e.what() << '\n';
+    return false;
+  }
+}
+
+bool MotionStrategy::LoadStrategy(const std::string& name,
+                                  const std::string& path) {
+  std::string filename(path);
+  filename.append(name);
+  filename.append(".json");
+  try {
+    std::ifstream file(filename);
+    if (!file.good()) {
+      std::printf("\nFile %s does not exist", filename.c_str());
+      return false;
+    }
+    contact_keyframes_.clear();
+    json json_keyframes = json::parse(file);
+    ContactKeyframe kf;
+    for (auto& json_kf: json_keyframes) {
+      from_json(json_kf, kf);
+      contact_keyframes_.push_back(kf);
+    }
+
+    current_keyframe_index_ = 0;
+    std::printf("\nKeyframe sequence loaded from %s", filename.c_str());
+    return true;
+  } catch (const std::exception &e) {
+    std::cerr << e.what() << '\n';
+    return false;
+  }
+}
+
+void to_json(json& j, const MotionStrategy& strategy) {
+  j = json{{"contact_keyframes", strategy.GetContactKeyframes()}};
+}
+
+void from_json(const json& j, MotionStrategy& strategy) {
+  strategy.SetContactKeyframes(
+      j.at("contact_keyframes").get<std::vector<ContactKeyframe>>());
+}
+
 }  // namespace mjpc::humanoid

mjpc/tasks/humanoid/interact/motion_strategy.h

@@ -17,13 +17,20 @@
 
 #include <mujoco/mujoco.h>
 
+#include <filesystem>
+#include <fstream>
 #include <string>
 #include <vector>
 
 #include "mjpc/tasks/humanoid/interact/contact_keyframe.h"
+#include "nlohmann/json.hpp"
+using json = nlohmann::json;
 
 namespace mjpc::humanoid {
 
+constexpr char kMotionStrategyFilePath[] =
+    SOURCE_DIR "/mjpc/tasks/humanoid/interact/strategies/";
+
 /*
 This class holds the motion strategy, e.g. given a sequence of keyframes, it
 manages the initial and current state and any changes to the sequence.
@@ -69,6 +76,10 @@ class MotionStrategy {
   void SetContactKeyframes(const std::vector<ContactKeyframe>& keyframes) {
     contact_keyframes_ = keyframes;
   }
+  bool SaveStrategy(const std::string& name,
+                    const std::string& path = kMotionStrategyFilePath);
+  bool LoadStrategy(const std::string& name,
+                    const std::string& path = kMotionStrategyFilePath);
 
   int NextKeyframe();
   void ClearKeyframes();
@@ -84,6 +95,8 @@ class MotionStrategy {
   mjtNum current_keyframe_success_start_time_;
 };
 
+void to_json(json& j, const MotionStrategy& strategy);
+void from_json(const json& j, MotionStrategy& strategy);
 }  // namespace mjpc::humanoid
 
 #endif  // MJPC_TASKS_HUMANOID_INTERACT_MOTION_STRATEGY_H_